Compute Aligned Training: Optimizing for Test Time Inference

1. Probleemstelling: Het misalignment tussen training en inferentie

Het artikel adresseert een fundamentele disconnectie in de ontwikkeling van moderne Large Language Models (LLM): trainingsdoelen zijn geoptimaliseerd voor prestaties op een enkel monster, terwijl de implementatie steeds meer steunt op strategieën voor compute tijdens de testfase (search).

• De Status Quo: Standaard post-training methoden zoals Supervised Fine-Tuning (SFT) en Reinforcement Learning (RL) optimaliseren voor de waarschijnlijkheid van een enkel correct antwoord (bijvoorbeeld het minimaliseren van Cross-Entropy of het maximaliseren van de verwachte beloning van één monster).
• Het Conflict: Moderne inferentiestrategieën (zoals Pass@N, Majority Vote, Best-of-N) genereren meerdere kandidaat-oplossingen en aggregeren deze om het beste antwoord te vinden.
  • Voorbeeld (Pass@N): Als een model een kans van 50% heeft om correct te zijn, levert het verhogen daarvan naar 90% afnemende meerwaarde op voor een Pass@100-metriek (waar succes al bijna zeker is). Standaard training duwt echter voortdurend voor een waarschijnlijkheid van 100%, wat capaciteit verspillen aan "makkelijke" monsters.
  • Voorbeeld (Majority Vote): Een model heeft slechts een meerderheid (bijvoorbeeld 35% waarschijnlijkheid) nodig om te winnen van concurrenten. Standaard training duwt voor $p \to 1.0$, wat "landslide"-overwinningen aanmoedigt in plaats van robuust consensus, wat de prestaties kan verslechteren wanneer het model gedwongen wordt te fungeren als ensemble.
• Gevolg: Dit misalignment veroorzaakt policy collapse (verminderde diversiteit) en verhindert dat modellen effectief gebruikmaken van verhoogde compute-budgetten tijdens de testfase.

2. Methodologie: Compute Aligned Training (CAT)

De auteurs stellen Compute Aligned Training (CAT) voor, een raamwerk dat trainingsdoelen afstemt op de specifieke inferentiestrategie die tijdens de testfase wordt gebruikt.

Kernconceptualisatie

De auteurs conceptualiseren inferentiestrategieën niet als a posteriori heuristieken, maar als differentieerbare operatoren ($T$) die het basisbeleid $\pi_\theta$ transformeren naar een effectieve verdeling tijdens de testfase $\tilde{\pi}_\theta$:
$$\tilde{\pi}_\theta(y|x) = T(\pi_\theta(\cdot|x), \phi)$$
waarbij $\phi$ hyperparameters vertegenwoordigt (bijvoorbeeld het monsterbudget $N$).

De Unificatie Verliesfunctie

In plaats van de negatieve log-waarschijnlijkheid van de grondwaarheid te minimaliseren onder het basisbeleid ($\pi_\theta$), minimaliseert CAT het verlies onder het effectieve systeemniveau-beleid ($\tilde{\pi}_\theta$):
$$L_{CAT}(\theta) = -\mathbb{E}_{(x, y^*)} [\log \tilde{\pi}_\theta(y^*|x)]$$

Gradiënt-Schaalmechanisme

Om dit computationeel haalbaar te maken zonder de volledige outputruimte te bemonsteren, leiden de auteurs een diagonale gradiëntbenadering af. Zij tonen aan dat de gradiënt van het CAT-doel kan worden benaderd door de standaardgradiënt te schalen met een wegingsfactor $w(p, \phi)$:
$$\nabla_\theta L_{CAT} \approx w(p, \phi) \cdot \nabla_\theta L_{CE}$$
Hierbij is $p$ de basiswaarschijnlijkheid van het juiste antwoord, en $w$ vertegenwoordigt de marginale nut van het verhogen van $p$ met betrekking tot de strategie tijdens de testfase.

Afgeleide Belangrijke Schalingsfactoren:

1. Pass@N (SFT & RL):
   • $w(p, N) = \frac{N p (1-p)^{N-1}}{1 - (1-p)^N}$
   • Gedrag: Onderdrukt gradiënten voor makkelijke monsters (hoge $p$) waar succes al gegarandeerd is, en richt het leren op "moeilijke" monsters waar marginale winsten de dekking verbeteren.
2. Majority Vote (SFT & RL):
   • $w(p, N, k) \approx \frac{k \binom{N}{k} p^k (1-p)^{N-k}}{\sum_{i=k}^N \binom{N}{i} p^i (1-p)^{N-i}}$
   • Gedrag: Concentreert gradiëntmassa dicht bij de consensusdrempel ($p \approx k/N$). Het negeert monsters die al winnen of hopeloos verliezen, en richt zich uitsluitend op het "tippunt".
3. Best-of-N (RL):
   • $w \propto N (P_{<y})^{N-1}$
   • Gedrag: Creëert "winner-take-all"-dynamiek. Gradiënten versterken alleen de monsters met de hoogste beloning in een batch, waardoor het model wordt aangemoedigd om naar de hoog-variatie, hoog-beloning staarten van de verdeling te zoeken in plaats van in te storten naar een veilig gemiddelde.

Theoretische Waarborgen

Het artikel bewijst dat voor Competitieve Strategieën (waarbij het verhogen van de waarschijnlijkheid van een rivaal de succeskans schaadt), de diagonale benadering een conservatieve onderschatting is van de ware gradiënt. Dit zorgt ervoor dat de optimizer nooit in de verkeerde richting stapt, en fungeert als een natuurlijke stabilisator tegen policy collapse.

3. Belangrijkste Bijdragen

1. Unificatie Operator Framework: Formaliseert strategieën tijdens de testfase als differentieerbare operatoren, en leidt een algemeen mechanisme af om training af te stemmen op willekeurige inferentiestrategieën (SFT en RL).
2. Extensie Buiten Pass@N: Generaliseert eerder werk (bijvoorbeeld Direct Coverage Optimization) naar Majority Vote en Best-of-N, en dekt zowel SFT- als RL-paradigma's.
3. Modaal Onafhankelijkheid: Demonstreert de toepasbaarheid van het raamwerk buiten LLM's om naar Proteïne Language Models (PLMs), en toont verbeterde prestaties in complexe, multi-modale beloningslandschappen.
4. Variatieanalyse: Biedt een theoretische en empirische analyse van inflatie van gradiëntvariatie, en stelt normalisatietechnieken voor (bijvoorbeeld Log-Gewogen schatters) om training te stabiliseren.

4. Experimentele Resultaten

De auteurs evalueerden CAT op de MATH-benchmark (met Mistral-7B) en Proteïne Generatie (met ProtGPT2).

A. Taalmodellen (MATH Benchmark)

• Pass@N (SFT): Modellen getraind met CAT voor $N=64$ offerden enkel-monster nauwkeurigheid op (Pass@1 daalde met ~2,3%), maar behaalden een absolute verbetering van 7,8% in Pass@64-dekking vergeleken met standaard SFT.
• Majority Vote (SFT): CAT-modellen presteerden consistent beter dan baselines over alle inferentiebudgetten. Er werd een "sweet spot" waargenomen bij $N=16$, waarbij het model bias en variantie optimaal in evenwicht bracht.
• Reinforcement Learning (GRPO): Standaard RL leed aan policy collapse, en verzadigde snel (Pass@32 $\approx$ 35,8%). CAT-getrainde modellen behielden de nodige variantie voor search, en behaalden 40,0% bij Pass@32.
• Warmup-effect: Het gebruik van CAT voor de SFT-warmup-fase verbeterde de daaropvolgende RL-prestaties aanzienlijk, en verhoogde zelfs de enkel-shot (Pass@1) nauwkeurigheid door de redeneringspaden van het model te verrijken.

B. Proteïne Generatie (Proteïne Language Models)

• Ondoorwaarde (Valley of Death): In een multi-modaal beloningslandschap (lokaal optimum bij $h \approx 0.35$, globaal optimum bij $h \approx 0.75$, gescheiden door een "vallei" van aggregatie), bleef standaard RL steken in het lokale optimum. Best-of-N CAT slaagde erin de vallei succesvol te "overslaan", en vond het globale maximum door een hoge distributievariantie te behouden.
• Conditionele Generatie: CAT-modellen leerden sequenties te genereren met inverse hydrofobiciteit ten opzichte van de invoer, en toonden superieure schaling en afstemming op de relationele taak in vergelijking met standaard RL.

5. Betekenis en Implicaties

• Paradigmaverschuiving: Het artikel betoogt dat trainingsdoelen moeten worden ontworpen met het downstream inferentieproces in gedachten. Aangezien compute tijdens de testfase de primaire hefboom wordt voor het schalen van modelcapaciteit, moet training optimaliseren voor het geaggregeerde resultaat, niet voor het enkele monster.
• Efficiëntie: CAT bereikt deze winsten zonder de trainingscomputekosten te verhogen. Het weegt bestaande gradiënten simpelweg opnieuw op basis van de marginale nut van het monster voor de strategie tijdens de testfase.
• Stabiliteit versus Exploratie: Het raamwerk biedt een principiële manier om enkel-monster precisie in te leveren voor distributie-diversiteit, waardoor modellen effectief search-procedures zoals Best-of-N en Majority Vote kunnen gebruiken zonder in mode-zoekend gedrag in te storten.
• Generaliseerbaarheid: Door de effectiviteit aan te tonen in zowel LLM's als Proteïne LM's, suggereert het artikel dat "Compute Aligned Training" een universeel principe is voor elk generatief model dat search tijdens de testfase gebruikt.

Kortom, Compute Aligned Training overbrugt de kloof tussen leren en search, en zorgt ervoor dat modellen niet alleen "slim" zijn in het genereren van één antwoord, maar geoptimaliseerd zijn om "slim" te zijn in het vinden van het beste antwoord wanneer het computationele budget voor search beschikbaar is.