GradCuit: Credit-Assigned Gradient Flow Enables Robust and Interpretable Test-Time Latent Reasoning

Technische Samenvatting: GradCuit

Probleemstelling

Optimalisatie-gebaseerd latent redeneren beoogt de outputs van Large Language Models (LLM's) te verbeteren door tijdens de testtijd specifieke continue latente toestanden te optimaliseren terwijl de modelparameters bevroren blijven. Echter, bestaande methoden (bijv. LatentSeek, LTPO, MILR) kampen met een fundamenteel credit-assignment probleem (toeschrijvingsprobleem). Deze benaderingen verbinden latente variabelen doorgaans met het redeneertraject via gedecodeerde tokens. Dit creëert een informatiebottleneck waarbij optimalisatiesignalen door het discrete decodeerproces moeten terugpropageren om de latente variabelen te bereiken. Dit leidt tot:

1. Indirecte Optimalisatie: De relatie tussen latente updates en downstream redeneertokens is verstrengeld en indirect.
2. Ondoorzichtige Dynamiek: Het is moeilijk te interpreteren hoe specifieke latente variabelen de daaropvolgende voorspellingen of het uiteindelijke antwoord beïnvloeden, aangezien het discrete generatiepad de interne dynamiek vertroebelt.

Methodologie: GradCuit

De auteurs stellen GradCuit (Gradient through Circuit) voor, een methode die het self-attention mechanisme van de Transformer gebruikt als een computationeel circuit om directe credit assignment mogelijk te maken.

Kernmechanisme

In plaats van latente variabelen te definiëren bij de output van het model (zoals bij LatentSeek), voegt GradCuit een kleine set optimaliseerbare latente toestanden ($z^{(l)}$) direct in de intermediaire verborgen toestandruimte van een geselecteerde Transformer-laag ($l$) in.

1. Invoeging: Voor een prompt $c$ en eerder gegenereerde tokens $x_{<t}$, berekent het model de verborgen representaties tot laag $l$. De optimaliseerbare latente variabelen worden samengevoegd met deze representaties: $[h^{(l)}_c, z^{(l)}, h^{(l)}_{x<t}]$.
2. Forward Flow: Deze samengevoegde sequentie gaat door de resterende Transformer-lagen ($l+1$ tot $M$) en de taalmodel-kop om de volgende token-distributie te genereren.
3. Backward Flow (Credit Assignment): Omdat de latente variabelen deelnemen aan het self-attention mechanisme naast de token-representaties, heeft elke gegenereerde token's log-waarschijnlijkheid een differentieerbaar pad naar elke voorafgaande latente staat via de resterende lagen.
4. Optimalisatie: De methode optimaliseert de latente variabelen met een policy-gradient-achtige update:
   $$z^{(l)} \leftarrow z^{(l)} + \eta \nabla_{z^{(l)}} J(z^{(l)})$$
   waarbij de gradiënt $J$ de beloningsgewogen bijdragen van alle gegenereerde tokens aggregeert:
   $$\nabla_{z^{(l)}_i} J = \sum_{t=1}^T \mathbb{E}_{x \sim \pi} [R(x, c) \times \nabla_{z^{(l)}_i} \log \pi(x_t | x_{<t}, z^{(l)}, c)]$$
   Dit stelt de methode in staat om beloningssignalen van de gehele voortzetting direct toe te wijzen aan de latente variabelen zonder terug te propageren door het discrete token-decodeerproces.

Implementatiedetails

• Latente Constructie: Een vaste tekstuele prefix (bijv. "Let's think about this problem...") wordt gerenderd, en een trainbare offset $\Delta z^{(l)}$ wordt toegevoegd aan de verborgen representaties bij laag $l$.
• Zelf-beloning (Self-Reward): De LLM zelf fungeert als verifieerder en geeft een binaire beloning ($0$ voor correct, $-1$ voor incorrect) op basis van het geëxtraheerde uiteindelijke antwoord.
• Optimalisatie-loop: Het model genereert greedily een respons, evalueert deze, en werkt $\Delta z^{(l)}$ bij via Adam. Dit herhaalt zich totdat het antwoord correct is of een budget (maximaal 10 updates) is bereikt.

Belangrijkste Bijdragen

1. Circuit-achtige Latente Redenering: GradCuit introduceert een methode om leerbare latente toestanden in te voegen op een intermediaire Transformer-laag, waarbij gebruik wordt gemaakt van self-attention als een direct computationeel pad voor zowel forward interactie als backward credit assignment.
2. Verbeterde Robuustheid: Door beloningsgewogen gradiënten direct van gegenereerde tokens naar latente toestanden te propageren, bereikt GradCuit een grotere optimalisatiestabiliteit vergeleken met eerdere latente redeneermethoden.
3. Interpreteerbaarheid: Het circuit-achtige gradiëntpad maakt inspectie van latente dynamiek mogelijk. De auteurs tonen aan dat de invloed van latente variabelen zich concentreert op "redeneer-connector" tokens (bijv. "omdat", "daarom"), wat onthult hoe geoptimaliseerde toestanden het redeneerproces sturen.

Experimentele Resultaten

De auteurs hebben GradCuit geëvalueerd over vijf instruction-tuned backbones (LLaMA-3.2-3B, LaMA-3.1-8B, Qwen2.5-7B, Qwen2.5-14B, Qwen3-4B) en drie redeneer-benchmarks (GPQA-Diamond, GSM8K, MATH-500).

• Prestaties: GradCuit behaalde een gemiddelde nauwkeurigheid van 64,5%, waarmee het Chain-of-Thought (CoT) met 6,6 procentpunt versloeg en de sterkste concurrerende methode (LatentSeek) met 2,4 punten. Het behaalde het beste individuele resultaat in 23 van de 30 settings.
• Robuustheid tegen Hyperparameters: Over zeven instellingen voor de leersnelheid behield GradCuit een hoge prestatie met een standaarddeviatie van 0,82, wat aanzienlijk lager is dan de 1,53 van LatentSeek.
• Robuustheid tegen Optimalisatierichting: Zelfs wanneer beloningsafgeleide gradiënten werden vervangen door willekeurige Gaussische richtingen (een random-walk variant), bleef GradCuit competitief met LatentSeek, wat suggereert dat de latente ruimte inherent toegankelijker en robuuster is.
• Laag-gevoeligheid: Optimale prestaties werden waargenomen wanneer de latente variabelen werden ingevoegd bij 25% tot 50% van de netwerkdiepte, wat een balans biedt tussen contextualisering en voldoende downstream computation voor verfijning.

Betekenis

Het artikel beweert dat GradCuit een nieuwe as opent voor robuuste en interpreteerbare test-time scaling. Door direct de interne redeneertoestanden te optimaliseren vanuit outcome-feedback, stelt de methode LLM's in staat om "aan te passen hoe ze redeneren" in plaats van enkel outputs te regenereren, samplen of herrangschikken. Het vermogen om direct credit toe te wijzen via het attention-circuit van de Transformer lost de ondoorzichtigheid van eerdere latente redeneermethoden op, en biedt een pad naar stabielere en begrijpelijke test-time adaptatie zonder de basisparameters van het model bij te werken.