Technische Samenvatting: Gezamenlijke Training van Multi-Token Predictie in Versterkend Leren via Optimale Coëfficiëntkalibratie

Probleemstelling

Versterkend Leren vanuit Verifieerbare Beloningen (RLVR) is het standaardparadigma geworden voor het verbeteren van de redeneercapaciteiten van Large Language Models (LLM's). Tegelijkertijd heeft Multi-Token Predictie (MTP), een hulpdoelwit dat tijdens pretraining wordt getraind om meerdere toekomstige tokens te voorspellen, aanzienlijke voordelen aangetoond in steekproeffICIëntie en als een conceptmodel voor speculatieve decoding.

Intuïtief zou het gezamenlijk optimaliseren van MTP en RL-doelwitten tijdens post-training modellen in staat moeten stellen te profiteren van multi-token supervisie terwijl ze redeneercapaciteiten verwerven. Empirisch bewijs toont echter aan dat gezamenlijke MTP-RL-training vaak leidt tot significante prestatiedegradering van het hoofdmodel. Bijgevolg gebruiken huidige praktijken (bijvoorbeeld in frameworks zoals veRL en slime) doorgaans het loskoppelen van MTP-gradiënten van het hoofdmodel tijdens RL-training. Deze loskoppeling dwingt MTP-parameters zich alleen aan te passen via onafhankelijke post-hoc fijnafstemming, waardoor de potentiële winsten van gelijktijdige multi-token supervisie worden opgeofferd. Het kernprobleem dat wordt aangepakt is: Waarom faalt gezamenlijke MTP-RL-training, en kan er een principiële strategie worden ontworpen om effectieve gezamenlijke training mogelijk te maken?

Methodologie

Theoretische Analyse

De auteurs benaderen het probleem vanuit een optimalisatieperspectief, waarbij ze het effect per stap van MTP op de RL-doelwitfunctie $J(\theta)$ analyseren. Aannemende dat de doelwitfunctie $L$-glad is, leiden ze een ondergrens af voor de verbetering per stap. Deze decompositie onthult dat het effect van MTP wordt bepaald door twee concurrerende termen:

1. Eerste-orde Correlatieterm: Bepaald door de directionele overeenstemming (inwendig product) tussen de RL-gradiënt ($g_{RL}$) en de MTP-gradiënt ($g_{MTP}$).
2. Tweede-orde Perturbatieboete: Een niet-positieve term die de perturbatie weerspiegelt die wordt geïntroduceerd door de grootte van de MTP-gradiënt.

Prestatieverbetering treedt alleen op als de correlatieterm de boeteterm weegt.

Analyse van Bestaande Regimes

Met behulp van deze decompositie verenigt en verklaart het artikel drie bestaande trainingsregimes:

• Loskoppelen (Detach): $g_{MTP} = 0$ met betrekking tot het hoofdmodel. $\Delta_{MTP} = 0$. Dit is veilig, maar offert potentiële winsten op.
• Cross-Entropy (CE) Verlies: De correlatieterm is in verwachting bijna nul, omdat RL steekproeven met hoge beloning zwaarder weegt terwijl CE alle steekproeven gelijk behandelt. Bijgevolg domineert de negatieve tweede-orde boete, wat leidt tot degradatie.
• Beleidverlies (Policy Loss): Gebruikt dezelfde RL-doelwitformulering voor MTP. In het begin zijn de gradiënten uitgelijnd en domineert de correlatieterm, wat winsten oplevert. Echter, naarmate de training vordert en het model platte gebieden van het verlieslandschap binnengaat, vervaagt de correlatie terwijl de gradiëntnorm (en dus de boete) aanhoudt. Dit leidt tot een "stijg-then-val" prestatiecurve waarbij een vaste coëfficiënt de fase-overgang niet kan volgen.

Voorgestelde Oplossing: Optimale Coëfficiëntkalibratie (OCC)

Om de fase-overgang aan te pakken, stellen de auteurs Optimale Coëfficiëntkalibratie (OCC) voor, een adaptief schema dat de optimale MTP-coëfficiënt ($\lambda$) online bijhoudt.

1. Gesloten-vorm Optimum: De analyse levert een gesloten-vorm uitdrukking op voor de optimale $\lambda^*$ op elk moment, die evenredig is met de verhouding tussen de gradiëntcorrelatie ($c$) en het kwadraat van de gradiëntnorm ($v^2$).
2. Log-waarschijnlijkheidsproxy: Het berekenen van volledige modelgradiënten voor $c$ en $v^2$ is in grote schaal gedistribueerde systemen onbetaalbaar duur. OCC maakt gebruik van een log-waarschijnlijkheidsproxy afgeleid van de verandering in log-waarschijnlijkheid onder updates met kleine stappen ($\delta \approx \log \pi_\theta - \log \pi_{old}$).
   • Uitlijningsproxy ($\hat{c}_t$): Schat het inwendig product.
   • Variantieproxy ($\hat{v}^2_t$): Schat het kwadraat van de norm.
3. Adaptieve Update: De coëfficiënt wordt bijgewerkt als $\lambda_t = \lambda_+ \cdot \frac{\hat{c}_t}{\hat{v}^2_t + \epsilon}$. Hierdoor kan OCC de weging van MTP dynamisch aanpassen, waardoor positieve winsten gedurende de hele training worden behouden zonder volledige gradiënten te hoeven berekenen.

Belangrijkste Bijdragen

1. Theoretische Unificatie: Het artikel biedt een theoretische analyse die de regimes Loskoppelen, Cross-Entropy en Beleidverlies verenigt, en hun succes of falen verklaart op basis van het samenspel tussen eerste-orde correlatie en tweede-orde perturbatie.
2. OCC-algoritme: Het stelt Optimale Coëfficiëntkalibratie voor, een adaptieve procedure die de MTP-coëfficiënt online kalibreert met behulp van een log-waarschijnlijkheidsproxy. Deze methode vereist verwaarloosbare computerefficiëntie en vermijdt de berekening van volledige modelgradiënten.
3. Empirische Validatie: Uitgebreide experimenten over meerdere modellen (MiMo-7B, GLM-4.5-Air) en algoritmen (DAPO, GSPO) op wiskundige redeneerbenchmarks op competitieniveau tonen aan dat OCC consequent de Loskoppel-baseline evenaart of overtreft.

Experimentele Resultaten

De auteurs hebben OCC geëvalueerd op zes benchmarks: AIME24, AIME25, AMC, MATH, Minerva en OlympiadBench.

• Prestaties: Over alle benchmarks en modelschalen (7B dense en 106B MoE) heen overtrof OCC consequent de Loskoppel-baseline en het vaste-coëfficiënt Beleidverlies.
  • Op MiMo-7B met DAPO behaalde OCC een gemiddelde nauwkeurigheid van 61,7%, wat beter was dan Loskoppelen (58,9%) en Beleidverlies (57,4%).
  • Op GLM-4.5-Air behaalde OCC een gemiddelde nauwkeurigheid van 67,6% versus 65,9% voor Loskoppelen.
• Trainingsdynamiek: Waar Beleidverlies een stijg-then-val trend vertoonde (degraderend in latere stadia), behield OCC stabiele verbeteringen gedurende de hele training. CE-verlies presteerde consequent slechter dan Loskoppelen.
• Efficiëntie: OCC introduceerde verwaarloosbare overhead. De wandkloktijd voor training per stap voor OCC (8,07s) was vergelijkbaar met Loskoppelen (8,11s), terwijl volledige model-backpropagatie 45,06s vereiste (een vertraging van 5,6×).
• Ablatie: Sensitiviteitsanalyse toonde aan dat de vooraf gedefinieerde schalingsfactor $\lambda_+$ robuust is over een reeks waarden, waarbij $\lambda_+ = 1,0$ de beste gemiddelde prestatie oplevert.

Betekenis en Claims

Het artikel claimt dat het falen van gezamenlijke MTP-RL-training niet inherent is, maar voortkomt uit een mismatch tussen statische coëfficiënten en de dynamische drift van gradiëntuitlijning tijdens training. Door een theoretische decompositie van het optimalisatielandschap te bieden en OCC voor te stellen, tonen de auteurs aan dat MTP veilig kan worden heringevoerd in RL post-training.

De betekenis ligt in het in staat stellen van modellen om gelijktijdig redeneercapaciteiten te verwerven via RL en multi-token representaties te benutten zonder prestaties op te offeren of onbetaalbare computerkosten te maken. Het werk suggereert dat adaptieve kalibratie essentieel is voor gezamenlijke training in niet-convexe deep learning-landschappen, en biedt een principiële alternatief voor de huidige industriestandaard van gradiëntloskoppeling.

Beperkingen

De auteurs merken op dat hun theoretische analyse berust op $L$-gladheidsaannames die mogelijk niet exact gelden in niet-convexe landschappen, waardoor ze kwalitatieve in plaats van kwantitatieve richtlijnen bieden. De log-waarschijnlijkheidsproxy is een eerste-orde benadering die onder grote leersnelheden zijn nauwkeurigheid kan verliezen. Bovendien zijn de experimenten gericht op wiskundig redeneren met verifieerbare beloningen; generalisatie naar settings zoals RLHF met geleerde beloningsmodellen of open-ended generatie moet nog worden gevalideerd.