Technische Samenvatting: Het blootleggen van latente redeneerstrategieën in taalmodellen

Probleemstelling

Pretrained taalmodellen (LM's) die getraind zijn op redeneertaken, lossen problemen vaak op met behulp van meerdere, verschillende, hoogwaardige strategieën (bijv. bewijs door contradictie versus inductie, of verschillende sorteeralgoritmen). Deze strategieën zijn echter impliciet en verstrengeld binnen de responsdistributie $p_\theta(y|x)$ van het model. Er is geen expliciete variabele om te identificeren of te controleren welke strategie wordt gebruikt tijdens de generatie.

Het artikel formuleert het probleem van het ontbinden van deze responsdistributie in een gestructureerde, strategie-geconditioneerde representatie. Het doel is om een router-generator factorisatie $(r_\phi(z|x), g_\phi(y|x, z))$ te leren waarvoor geldt:

1. Marginale Getrouwheid (Fidelity): De geïntegreerde distributie $\int r_\phi(z|x)g_\phi(y|x, z) dz$ benadert de oorspronkelijke basismodel-distributie $p_\theta(y|x)$.
2. Strategie-alignement: De latente variabele $z$ is informatief en strategie-gealigneerd, wat betekent dat het variëren van $z$ de gebruikte hoogwaardige strategie verandert, en dat dezelfde latente waarde een consistente strategische betekenis behoudt over verschillende inputs.

De kernuitdaging: In tegenstelling tot standaard latente variabele modellering waarbij een model vanaf nul wordt getraind, bevat de basismodel $p_\theta$ al de doeldistributie zonder enige latente variabele. Daarom is een factorisatie die $z$ negeert (waarbij $g_\phi(y|x, z) = p_\theta(y|x)$ voor alle $z$) beschikbaar bij initialisatie en is deze globaal optimaal voor doelstellingen die alleen de marginale fit meten. Standaard variationele inferentie (bijv. conditionele ELBO) lijdt in deze setting aan posterior collapse, waarbij er geen stimulans is voor $z$ om de variatie op strategieniveau te coderen.

Methodologie

Architectuur

De auteurs passen een pretrained Transformer aan naar een router-generator framework:

• Router ($r_\phi$): Leest de input $x$ en geeft een distributie over een continue latente strategie $z$ (gemodelleerd als een diagonale Gaussische distributie).
• Generator ($g_\phi$): Een autoregressief model dat geconditioneerd is op $x$ en een geprojecteerde embedding $E(z)$ van de gesamplede latente waarde.
• Posterior ($q_\xi$): Een apart model dat alleen tijdens de training wordt gebruikt om $z$ af te leiden uit de volledige input-respons paren $(x, y)$.
• Parameterisatie: De router en generator delen parameters die geïnitialiseerd zijn vanuit het basismodel $\theta$, en worden bijgewerkt via LoRA-stijl low-rank adaptatie.

Model-gestuurde Reconstructie-doelstelling

Om de ondergespecificeerdheid van de standaard ELBO te overwinnen, stellen de auteurs een model-gestuurde reconstructie-doelstelling voor. Deze doelstelling gebruikt het vaste basismodel $p_\theta$ als referentie om te sturen waar de latente variabele $z$ de variatie moet verklaren.

De doelstelling $J_\theta$ bestaat uit een reconstructieterm $R_\theta$ en een KL-strafterm:
$$J_\theta(x, y; \phi, \xi) = R_\theta(x, y; \phi, \xi) + \beta \text{KL}(q_\xi(\cdot|x, y) \parallel r_\phi(\cdot|x))$$

Belangrijke componenten van de reconstructieterm:

1. Normalisatie door het Basismodel: De reconstructieverlies wordt genormaliseerd door de verwachte responsverlies van het basismodel ($c_\theta$), waardoor de loss wordt uitgedrukt in "eenheden van het basismodel-referentiekader".
2. Surprisal-gestuurde Weging: In plaats van uniforme token-weging, wordt de loss gewogen door de token-niveau surprisal van het basismodel ($b_{\theta,t}$).
   • $b_{\theta,t}(x, y) = -\log p_\theta(y_t | x, y_{<t})$
   • Tokens met een hoge surprisal (splitsingspunten waar het basismodel onzeker is) ontvangen hogere reconstructiedruk.
   • Dit stimuleert $z$ om variatie te coderen op posities waar het basismodel meerdere voortzettingen plausibel laat, wat typisch overeenkomt met keuzepunten voor strategieën.

De reconstructieterm wordt gedefinieerd als een gewogen, genormaliseerd deel van de loss van het basismodel:
$$R_\theta = \frac{1}{c_\theta} \mathbb{E}_{z \sim q_\xi} \left[ \sum_{t} a_{\theta,t}(x, y) [-\log g_\phi(y_t | x, z, y_{<t})] \right]$$
waarbij gewichten $a_{\theta,t}$ interpoleren tussen uniforme druk en surprisal-gestuurde druk.

Experimentele Opstelling en Benchmark

De auteurs introduceren een gecontroleerde benchmark van multi-strategische algoritmische taken om strategie-alignement meetbaar te maken.

• Taken: Lijst-sommatie, sorteeralgoritmen, grid-padvinden, lineaire vergelijkingen oplossen, basisconversie en meercijferige optelling.
• Ontwerp: Elke taak staat meerdere geldige oplossingstrajecten (strategieën) toe voor dezelfde probleeminstantie. Grondwaarheids-strategielabels worden gegenereerd maar tijdens de training achtergehouden, en worden enkel gebruikt voor evaluatie.
• Evaluatiemetrieken:
  1. Distributionele Getrouwheid: Het deel van de gegenereerde trajecten dat geldige oplossingen voor de taak blijft.
  2. Strategie-alignement (Analoge Consistentie): De waarschijnlijkheid dat het samplen van een latente $z$ uit een broninput $X$ en het toepassen ervan op een gerelateerde doelinput $X'$ resulteert in een oplossingstraject $Y'$ die dezelfde strategie gebruikt als $Y$ gegenereerd uit $X$.

Belangrijkste Resultaten

1. Impliciete Informatie Bestaat: Lineaire probes op de verborgen toestanden van het basismodel kunnen strategie-labels decoderen, wat bevestigt dat de doelstructuur impliciet aanwezig is vóór de factorisatie.
2. Falen van Standaard ELBO: Standaard variationele doelstellingen behouden de distributionele getrouwheid, maar resulteren in latent collapse, waarbij $z$ inert blijft en faalt om te aligneren met strategieën.
3. Succes van Model-gestuurde Reconstructie: De voorgestelde doelstelling slaagt erin latente codes te herstellen die:
   • Een hoge distributionele getrouwheid behouden (het behoud van de geldige outputs van het basismodel).
   • Een hoge Analoge Consistentie bereiken, wat aantoont dat $z$ herbruikbare strategische semantiek over verschillende inputs draagt.
   • Zich over de training heen splitsen in duidelijke taak- en strategie-specifieke regio's in de latente ruimte.
4. Ablatie-studies:
   • Het verwijderen van de globale schaal (normalisatie) of token-weging (surprisal-richting) verslechtert de strategie-alignement.
   • Het gebruik van inverse-surprisal weging (nadruk op tokens met lage surprisal) verzwakt de alignement, wat bevestigt dat druk moet worden uitgeoefend op onzekere/splitsende tokens om strategie-variatie te herstellen.

Theoretische Analyse

Het artikel initieert een theoretische analyse van het herstelprobleem. Het frameert de doelstelling als het maximaliseren van een geprofileerde score bestaande uit de wederzijdse informatie $I(Y; Z|X)$ minus een benaderingscomplexiteit-straf. De auteurs bieden voldoende voorwaarden waaronder een complexiteit-gestrafte variationele doelstelling een latente factorisatie selecteert die is uitgelijnd met een doelstrategiestructuur $S^*$, mits de "herstel-discrepantie" tussen de optimale oplossing en suboptimale oplossingen positief is.

Betekenis en Claims

Het artikel beweert een specifieke pathologische vorm van posterior collapse aan te pakken die uniek is voor het aanpassen van pretrained modellen in plaats van het trainen vanaf nul. Door een model-gestuurde reconstructie-doelstelling te introduceren, demonstreren de auteurs een methode om:

• De impliciete redeneerstrategieën te ontwarren in een expliciete, controleerbare latente variabele.
• De responsdistributie van het basismodel te behouden terwijl de latente variabele informatief wordt gemaakt.
• Interventie mogelijk te maken: De herstelde latente structuur maakt het mogelijk om generatie te sturen naar specifieke hoogwaardige strategieën (bijv. een specifief bewijsmechanisme afdwingen) zonder het basismodel opnieuw te trainen.

De auteurs merken op dat de methodologie gevalideerd is op synthetische algoritmische taken met grondwaarheids-strategieën. Ze geven expliciet aan dat het uitbreiden hiervan naar open-ended redeneertaken en het volledig karakteriseren van de verbinding tussen de theoretische oracle-selectiecriterium en de werkelijke trainingsdynamiek openstaande vragen zijn voor toekomstig werk. De genoemde potentiële toepassingen omvatten model-interpreteerbaarheid, controleerbaarheid en gerichte exploratie in reinforcement learning post-training.