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

Technische Samenvatting: Leren wat te voorspellen: Downstream-gestuurd taakontwerp voor voortgezette pretraining

1. Probleemstelling

Voortgezette pretraining (continued pretraining) is een standaardmethode om fundamentele modellen aan te passen aan nieuwe domeinen (bijv. wiskundige redenering of dichte visuele voorspelling). De huidige praktijk leunt echter op een grove feedbackloop. De optimalisatie-eenheid is een enkele zelfgesuperviseerde update (bijv. next-token prediction of gemaskeerde reconstructie), maar de selectie-eenheid is de gehele pretraining-run, geëvalueerd op downstream-benchmarks.

Praktijkgebruikers evalueren checkpoints, passen handmatig datamengsels of objectieven aan en starten runs opnieuw op. Dit proces maakt individuele pretraining-updates "blind" voor de vraag of ze bijdragen aan specifieke downstream-capaciteiten. Hoewel fine-tuning (SFT, RL) downstream-labels gebruikt om de leerder direct bij te werken, vervangt dit de zelfgesuperviseerde update in plaats van deze te sturen. De paper stelt de vraag: Kan een kleine set verifieerbare downstream-voorbeelden stap-niveau feedback bieden tijdens voortgezette pretraining om doelgerichte capaciteiten te verbeteren zonder de leerder direct te superviseren of de generalisatie te schaden?

2. Methodologie: V-pretraining

De auteurs introduceren V-pretraining (Value-based pretraining), een framework dat het leerproces splitst in twee componenten:

1. De Leerder (The Learner): Uitsluitend getraind via een zelfgesuperviseerde loss op een ongelabelde stroom. De parameters ($\theta$) worden nooit bijgewerkt door downstream-gradiënten.
2. De Taakontwerper (The Task Designer): Een lichtgewicht module met parameters ($\phi$) die dynamisch de zelfgesuperviseerde taak construeert (bijv. het selecteren van targets of views) voor de ongelabelde batch.

Kernmechanisme:
De ontwerper gebruikt een kleine feedbackset ($D_{fb}$) die downstream-labels bevat om de "waarde" van een kandidaat zelfgesuperviseerde update te schatten.

• Gradiënt-uitlijning: Laat $g_{pre}$ de gradiënt zijn van de voorgestelde zelfgesuperviseerde update en $g_{down}$ de gradiënt van de downstream-loss op de feedbackset.
• Waarde-objectief: De paper stelt dat de eerste-orde verandering in de downstream-loss benaderd kan worden door het inproduct $V = g_{down}^\top g_{pre}$.
• Trainingslus:
  1. De ontwerper construeert een taak (bijv. een zachte target-distributie of een specifieke view) voor de huidige ongelabelde batch.
  2. De ontwerper wordt bijgewerkt om $V$ te maximaliseren (d.w.z. om de zelfgesuperviseerde gradiënt uit te lijnen met de richting die de downstream-loss vermindert).
  3. De geconstrueerde targets of views worden losgekoppeld (stop-gradient toegepast).
  4. De leerder wordt bijgewerkt met behulp van de standaard zelfgesuperviseerde loss met deze losgekoppelde targets.

Instanties:

• Taal: De ontwerper leert adaptieve top-K soft targets voor next-token prediction. In plaats van een harde one-hot target, geeft de ontwerper een waarschijnlijkheidsdistributie uit over een kleine kandidaat-set (inclusclusief de ware token en top-$K-1$ modelvoorspellingen), gewogen door een gating-constante.
• Visie: De ontwerper leert instance-wise views of masks voor DINO-stijl zelfgesuperviseerd leren, waarbij de input-crops of augmentaties worden aangepast om beter aan te sluiten bij downstream dense prediction-taken (segmentatie, diepte-inschatting).

3. Belangrijkste Bijdragen

• Conceptueel: De paper herformuleert downstream-gestuurde voortgezette pretraining als leerbare taakconstructie. Het introduceert een kanaal voor sturing dat indirect is naar de leerder (via taakontwerp) maar direct naar de ontwerper, waardoor het vervangen van zelfgesuperviseerde updates door gesuperviseerde updates wordt vermeden.
• Algoritmisch: Het stelt een computationeel efficiënt, stap-niveau objectief voor ($g_{down}^\top g_{pre}$) dat dient als een eerste-orde surrogaat voor downstream-verbetering. Dit voorkomt de onhandelbaarheid van het differentiëren door volledige pretraining-trajecten (bilevel optimalisatie).
• Empirisch: Uitgebreide, qua wandkloktijd gematchte experimenten tonen aan dat V-pretraining de doelgerichte capaciteiten verbetert zonder de algemene representaties te laten instorten.

4. Experimentele Resultaten

De auteurs voerden compute-gematchte experimenten uit (dezelfde wandkloktijd, hardware en initiële condities) over taal- en visiemodaliteiten.

Taal (Wiskundig redeneren):

• Setup: Voortgezette pretraining van Qwen-modellen op NuminaMath-CoT met gebruik van 1.024 GSM8K-voorbeelden als feedback.
• Resultaten: V-pretraining verbeterde de GSM8K Pass@1 aanzienlijk.
  • Qwen2.5-0.5B: Enkele-run winst van +7,4 punten (22,20 $\to$ 29,60), een relatieve verbetering van 33%.
  • Qwen1.5-4B: Verbeterde van 56,48 naar 58,98.
• Generalisatie: De methode heeft de bredere capaciteiten niet geschaad. MMLU-scores bleven stabiel voor 4B en 7B modellen, hoewel het 0,5B model een lichte daling vertoonde, wat suggereert dat kleine leerders gevoeliger zijn voor smalle feedback.

Visie (Dense Prediction):

• Setup: Voortgezette zelfgesuperviseerde training van DINOv3-backbones op ImageNet-1K met ADE20K (segmentatie) en NYUv2 (diepte) als feedback.
• Resultaten:
  • ADE20K: mIoU verbeterde van 51,33 naar 52,40 voor ViT-L.
  • NYUv2: RMSE daalde van 0,5752 naar 0,5522 voor ViT-L.
  • Generalisatie: De ImageNet lineaire classificatie-accuratesse werd behouden of licht verbeterd, wat aangeeft dat het model niet over-gespecialiseerd is geraakt voor de dense taken ten koste van globale herkenning.
  • Transfer: Instance retrieval op Oxford/Paris datasets liet verbeteringen zien in de meeste protocollen, wat bevestigt dat de representaties robuust bleven.

Ablaties en Controles:

• Decontaminatie: Het verwijderen van bijna-duplicaten van GSM8K/MATH uit de trainingsstroom elimineerde de winst niet, wat simpel geheugen (memorization) uitsluit.
• Shortcut Ruling: De winsten waren niet toe te schrijven aan generieke smoothing, zelf-distillatie of willekeurige ruis. Varianten die gebruikmaakten van willekeurige feedback of uniforme smoothing presteerden slechter dan de uitgelijnde V-pretraining.
• Overhead: De taakontwerper introduceert een reductie van ~9,4% in token-throughput en een toename van ~16,4% in piekgeheugen, maar deze kosten zijn meegenomen in de wandkloktijd-gematchte vergelijkingen.

5. Betekenis en Claims

De paper claimt dat V-pretraining een computationeel efficiënt complement biedt aan directe post-training methoden (zoals SFT of RL).

• Efficiëntie: Het maakt het mogelijk om signalen van downstream-capaciteit te injecteren tijdens de fase van voortgezette pretraining, wat potentieel de noodzaak voor dure downstream fine-tuning later vermindert.
• Mechanisme: Het demonstreert dat downstream-feedback de constructie van zelfgesuperviseerde taken kan sturen (wat het model voorspelt) in plaats van de leerder zelf (hoe het model wordt bijgewerkt).
• Bescheidenheid: De auteurs stellen expliciet dat het waarde-objectief een lokale, één-stap garantie biedt (het maximaliseren van een ondergrens op onmiddellijke downstream-verbetering) in plaats van een langetermijn-convergentietheorema. Ze erkennen dat de methode kleine modellen te veel in de richting van smalle taken kan sturen en dat het uitlijningssignaal ruisgevoelig kan zijn.

Samenvattend stelt V-pretraining vast dat een kleine set downstream-voorbeelden effectief het aspect van "wat te voorspellen" van voortgezette pretraining kan sturen, waardoor specifieke capaciteiten worden verbeterd terwijl de generalisatie-eigenschappen van het fundamentele model behouden blijven.