Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Value-Based Pre-Training with Downstream Feedback

Shuqi Ke, Giulia Fanti

TL;DR

Value-Based Pre-Training with Downstream Feedback (V-Pretraining) introduces a lightweight task designer that, trained on a small set of verifiable downstream tasks, reshapes the pretraining signal so that each unlabeled gradient step aligns with downstream value. The approach uses an influence-style first-order objective, $\mathcal{V}(\phi; \theta) = g_{down}(\theta)^\top g_{pre}(\theta; \phi)$, to update the designer while the learner remains trained on unlabeled data, effectively turning a large open-loop pretraining process into a controlled trajectory toward desired capabilities. Empirically, V-Pretraining yields tangible gains in language reasoning (GSM8K) and dense-vision tasks (ADE20K, NYUv2) under fixed compute budgets, with modest overhead and robust generalization across tasks and model sizes. The work highlights a scalable path to compute-efficient capability shaping, bridging pretraining and alignment, and suggests directions for handling non-differentiable feedback and stronger post-training integration.

Abstract

Can a small amount of verified goal information steer the expensive self-supervised pretraining of foundation models? Standard pretraining optimizes a fixed proxy objective (e.g., next-token prediction), which can misallocate compute away from downstream capabilities of interest. We introduce V-Pretraining: a value-based, modality-agnostic method for controlled continued pretraining in which a lightweight task designer reshapes the pretraining task to maximize the value of each gradient step. For example, consider self-supervised learning (SSL) with sample augmentation. The V-Pretraining task designer selects pretraining tasks (e.g., augmentations) for which the pretraining loss gradient is aligned with a gradient computed over a downstream task (e.g., image segmentation). This helps steer pretraining towards relevant downstream capabilities. Notably, the pretrained model is never updated on downstream task labels; they are used only to shape the pretraining task. Under matched learner update budgets, V-Pretraining of 0.5B--7B language models improves reasoning (GSM8K test Pass@1) by up to 18% relative over standard next-token prediction using only 12% of GSM8K training examples as feedback. In vision SSL, we improve the state-of-the-art results on ADE20K by up to 1.07 mIoU and reduce NYUv2 RMSE while improving ImageNet linear accuracy, and we provide pilot evidence of improved token efficiency in continued pretraining.

Value-Based Pre-Training with Downstream Feedback

TL;DR

Value-Based Pre-Training with Downstream Feedback (V-Pretraining) introduces a lightweight task designer that, trained on a small set of verifiable downstream tasks, reshapes the pretraining signal so that each unlabeled gradient step aligns with downstream value. The approach uses an influence-style first-order objective, , to update the designer while the learner remains trained on unlabeled data, effectively turning a large open-loop pretraining process into a controlled trajectory toward desired capabilities. Empirically, V-Pretraining yields tangible gains in language reasoning (GSM8K) and dense-vision tasks (ADE20K, NYUv2) under fixed compute budgets, with modest overhead and robust generalization across tasks and model sizes. The work highlights a scalable path to compute-efficient capability shaping, bridging pretraining and alignment, and suggests directions for handling non-differentiable feedback and stronger post-training integration.

Abstract

Can a small amount of verified goal information steer the expensive self-supervised pretraining of foundation models? Standard pretraining optimizes a fixed proxy objective (e.g., next-token prediction), which can misallocate compute away from downstream capabilities of interest. We introduce V-Pretraining: a value-based, modality-agnostic method for controlled continued pretraining in which a lightweight task designer reshapes the pretraining task to maximize the value of each gradient step. For example, consider self-supervised learning (SSL) with sample augmentation. The V-Pretraining task designer selects pretraining tasks (e.g., augmentations) for which the pretraining loss gradient is aligned with a gradient computed over a downstream task (e.g., image segmentation). This helps steer pretraining towards relevant downstream capabilities. Notably, the pretrained model is never updated on downstream task labels; they are used only to shape the pretraining task. Under matched learner update budgets, V-Pretraining of 0.5B--7B language models improves reasoning (GSM8K test Pass@1) by up to 18% relative over standard next-token prediction using only 12% of GSM8K training examples as feedback. In vision SSL, we improve the state-of-the-art results on ADE20K by up to 1.07 mIoU and reduce NYUv2 RMSE while improving ImageNet linear accuracy, and we provide pilot evidence of improved token efficiency in continued pretraining.
Paper Structure (54 sections, 3 theorems, 26 equations, 3 figures, 4 tables, 1 algorithm)

This paper contains 54 sections, 3 theorems, 26 equations, 3 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminaries and Related Work
  3. Pretraining as Predictive Learning
  4. Related Works
  5. Pretraining with Downstream Feedback
  6. Learning to Design Pretraining Tasks
  7. Value Function for Downstream Feedback
  8. Algorithmic Instantiations
  9. Theoretic Guarantee
  10. Experiments
  11. Setup
  12. Evaluation on Selected Downstream Tasks
  13. Feedback Effects on Generalization
  14. Scaling Weak-to-Strong Supervision
  15. Ablation Studies
  16. ...and 39 more sections

Key Result

Theorem 3.1

Let $\theta^{+}=\theta-\eta\, g_{\mathrm{pre}}(\theta;\phi)$ for step size $\eta>0$ and define $g_{\mathrm{down}}(\theta)=\nabla_\theta L_{\mathrm{down}}(\theta)$. Under eq:smoothness, if $L_{\mathrm{down}}$ is $L$-smooth,

Figures (3)

  • Figure 1: Value-Based Pretraining with Downstream Feedback. Today, the learner $\theta$ trains on unlabeled data using a proxy objective $L_{\mathrm{pre}}$, for a frozen pretraining task. In V-Pretraining, a small task designer $\phi$ is trained on a small feedback set of verifiable downstream tasks with predefined value functions, but never updates the learner on downstream labels. $\phi$ thus reshapes the pretraining target (or views) so that the induced SSL update aligns with downstream improvement, calculated via the value function. Relative to current pretraining methods, V-Pretraining adds the components in the left blue box.
  • Figure 2: Token efficiency and multi-objective control. Left: GSM8K test Pass@1 versus unlabeled tokens processed for Qwen1.5-4B under matched learner-step budgets. Right: Tradeoff between segmentation (mIoU) and depth estimation (1-RMSE) induced by varying feedback and task-designer hyperparameters.
  • Figure 3: Scaling feedback coverage and inference-time compute.

Theorems & Definitions (6)

  • Theorem 3.1: Value lower bounds one-step downstream improvement
  • Proposition 3.2: Value is the first-order surrogate of one step bilevel optimization
  • Lemma 3.3: Unbiased stochastic value under independent sampling
  • proof : Proof of \ref{['thm:value_descent']}
  • proof : Proof of \ref{['prop:one_step_bilevel']}
  • proof : Proof of \ref{['lem:unbiased_value']}