Learning What to Predict: Downstream-Guided Task Design for Continued Pretraining

Technical Summary: Learning What to Predict: Downstream-Guided Task Design for Continued Pretraining

1. Problem Statement

Continued pretraining is a standard method for adapting foundation models to new domains (e.g., mathematical reasoning or dense vision prediction). However, current practice relies on a coarse feedback loop. The optimization unit is a single self-supervised update (e.g., next-token prediction or masked reconstruction), but the selection unit is the entire pretraining run evaluated on downstream benchmarks.

Practitioners typically evaluate checkpoints, manually revise data mixtures or objectives, and restart runs. This process renders individual pretraining updates "blind" to whether they contribute to specific downstream capabilities. While fine-tuning (SFT, RL) uses downstream labels to directly update the learner, this replaces the self-supervised update rather than guiding it. The paper asks: Can a small set of verifiable downstream examples provide step-level feedback during continued pretraining to improve target capabilities without directly supervising the learner or harming generalization?

2. Methodology: V-pretraining

The authors introduce V-pretraining (Value-based pretraining), a framework that separates the learning process into two components:

1. The Learner: Trained exclusively via a self-supervised loss on an unlabeled stream. Its parameters ($\theta$) are never updated by downstream gradients.
2. The Task Designer: A lightweight module with parameters ($\phi$) that dynamically constructs the self-supervised task (e.g., selecting targets or views) for the unlabeled batch.

Core Mechanism:
The designer uses a small feedback set ($D_{fb}$) containing downstream labels to estimate the "value" of a candidate self-supervised update.

• Gradient Alignment: Let $g_{pre}$ be the gradient of the proposed self-supervised update and $g_{down}$ be the gradient of the downstream loss on the feedback set.
• Value Objective: The paper posits that the first-order change in downstream loss is approximated by the inner product $V = g_{down}^\top g_{pre}$.
• Training Loop:
  1. The designer constructs a task (e.g., a soft target distribution or a specific view) for the current unlabeled batch.
  2. The designer is updated to maximize $V$ (i.e., to align the self-supervised gradient with the direction that reduces downstream loss).
  3. The constructed targets or views are detached (stop-gradient applied).
  4. The learner is updated using the standard self-supervised loss with these detached targets.

Instantiations:

• Language: The designer learns adaptive top-K soft targets for next-token prediction. Instead of a hard one-hot target, the designer outputs a probability distribution over a small candidate set (including the true token and top-$K-1$ model predictions), weighted by a gating constant.
• Vision: The designer learns instance-wise views or masks for DINO-style self-supervised learning, modifying the input crops or augmentations to better align with downstream dense prediction tasks (segmentation, depth estimation).

3. Key Contributions

• Conceptual: The paper reframes downstream-guided continued pretraining as learnable task construction. It introduces a channel for guidance that is indirect-to-learner (via task design) but direct-to-designer, avoiding the replacement of self-supervised updates with supervised ones.
• Algorithmic: It proposes a computationally efficient, step-level objective ($g_{down}^\top g_{pre}$) that serves as a first-order surrogate for downstream improvement. This avoids the intractability of differentiating through full pretraining trajectories (bilevel optimization).
• Empirical: Extensive wall-clock-matched experiments demonstrate that V-pretraining improves target capabilities without collapsing general-purpose representations.

4. Experimental Results

The authors conducted compute-matched experiments (same wall-clock time, hardware, and initial conditions) across language and vision modalities.

Language (Math Reasoning):

• Setup: Continued pretraining of Qwen models on NuminaMath-CoT using 1,024 GSM8K examples as feedback.
• Results: V-pretraining significantly improved GSM8K Pass@1.
  • Qwen2.5-0.5B: Single-run gain of +7.4 points (22.20 $\to$ 29.60), a 33% relative improvement.
  • Qwen1.5-4B: Improved from 56.48 to 58.98.
• Generalization: The method did not harm broader capabilities. MMLU scores remained stable for 4B and 7B models, though the 0.5B model showed a slight drop, suggesting small learners are more sensitive to narrow feedback.

Vision (Dense Prediction):

• Setup: Continued self-supervised training of DINOv3 backbones on ImageNet-1K using ADE20K (segmentation) and NYUv2 (depth) as feedback.
• Results:
  • ADE20K: mIoU improved from 51.33 to 52.40 for ViT-L.
  • NYUv2: RMSE decreased from 0.5752 to 0.5522 for ViT-L.
  • Generalization: ImageNet linear classification accuracy was preserved or slightly improved, indicating the model did not over-specialize to the dense tasks at the expense of global recognition.
  • Transfer: Instance retrieval on Oxford/Paris datasets showed improvements in most protocols, confirming the representations remained robust.

Ablations and Controls:

• Decontamination: Removing near-duplicates of GSM8K/MATH from the training stream did not eliminate the gains, ruling out simple memorization.
• Shortcut Ruling: The gains were not due to generic smoothing, self-distillation, or random noise. Variants using random feedback or uniform smoothing performed worse than the aligned V-pretraining.
• Overhead: The task designer introduces a ~9.4% reduction in token throughput and a ~16.4% increase in peak memory, but these costs are accounted for in the wall-clock-matched comparisons.

5. Significance and Claims

The paper claims that V-pretraining offers a computationally efficient complement to direct post-training methods (like SFT or RL).

• Efficiency: It allows the injection of downstream capability signals during the continued pretraining phase, potentially reducing the need for expensive downstream fine-tuning later.
• Mechanism: It demonstrates that downstream feedback can steer the construction of self-supervised tasks (what the model predicts) rather than the learner itself (how the model updates).
• Modesty: The authors explicitly state that the value objective provides a local, one-step guarantee (maximizing a lower bound on immediate downstream improvement) rather than a long-horizon convergence theorem. They acknowledge that the method can over-steer small models toward narrow tasks and that the alignment signal can be noisy.

In summary, V-pretraining establishes that a small set of downstream examples can effectively guide the "what to predict" aspect of continued pretraining, improving specific capabilities while maintaining the generalization properties of the foundation model.