Technische Samenvatting: IG-Lens

Probleemstelling

Decoder-only transformers worden vaak geconceptualiseerd als iteratieve inferentieprocessen waarbij de residuële stroom de lopende schatting van het volgende token verfijnt. Een cruciale interpretatievraag is: op welke specifieke laag wordt de waarschijnlijkheid van een voorspeld token daadwerkelijk beslist?

Bestaande instrumenten voor laag-gewijze uitlezing (readout) slagen er niet in dit precies te beantwoorden in de waarschijnlijkheidsruimte:

• Logit Lens & Tuned Lens: Deze rapporteren het waarschijnlijkheidsniveau per laag, maar bieden geen additieve decompositie. Hun schattingen zijn vertekend en niet-monotoon over de diepte, wat betekent dat de som van de waarschijnlijkheden per laag geen wiskundige betekenis heeft.
• Direct Logit Attribution (DLA) & Embedding-methoden: Deze zijn additief, maar opereren in de logit-ruimte. Omdat de softmax-functie zeer niet-lineair is, vertaalt een additieve decompositie van logits zich niet naar een additieve decompositie van waarschijnlijkheden, wat juist de grootheid van praktische interesse is.
• Layer Conductance & Layer Integrated Gradients (IG): Deze methoden integreren gradiënten per laag, maar kennen elke laag toe aan zijn eigen onafhankelijke baseline. Hierdoor "telescoperen" de termen per laag niet (ze tellen niet op) tot de totale verandering in voorspellingswaarschijnlijkheid.

Methodologie: IG-Lens

De auteurs introduceren IG-Lens, een methode die een exacte additieve decompositie van waarschijnlijkheid biedt over transformer-lagen heen.

Kernconcept

IG-Lens past een telescopische toepassing van Integrated Gradients (IG) toe langs een enkel pad door de verborgen toestanden (hidden states) van een baseline naar de laatste laag.

1. Paddefinitie: De methode selecteert een reeks lagen $0 = L_0 < L_1 < \dots < L_k = n$ en definieert een pad van verborgen toestanden $h_{L_0} \to h_{L_1} \to \dots \to h_{final}$.
2. Head-Once Readout: Cruciaal is dat de readout-functie $f(h)$ wordt gedefinieerd als de uiteindelijke unembedding toegepast één keer op een enkele verborgen vector: $f(h) = \text{softmax}(\text{head}(\text{norm}(h)))_{y_t}$. Deze functie bevat geen aandacht (attention) of MLP; alle cross-token mixing wordt verondersteld te hebben plaatsgevonden in de forward pass die de losgekoppelde verborgen toestanden produceerde.
3. Telescopische Attributie: De bijdrage van laag $L_j$ wordt gedefinieerd als de segment-IG van $f$ langs de rechte lijn van $h_{L_{j-1}}$ naar $h_{L_j}$. Volgens de gradiëntstelling is het integraal van de gradiënt over een segment gelijk aan het verschil tussen de eindwaarden. Het sommeren van deze segmenten levert een telescopische som op:
   $$\sum_{j=1}^k IGL_j = f(h_{final}) - f(h_{base}) = \Delta p$$
   Dit garandeert dat de som van de attributies exact gelijk is aan de verandering in doelwaarschijnlijkheid, waarbij de softmax en LayerNorm worden geïntegreerd in plaats van lineair benaderd of bevroren.

Voorspellingsbewuste Schatter (Prediction-Aware Estimator)

Hoewel het integraal exact is, introduceren eindige-stappen schatters meestal discretisatiefouten. IG-Lens gebruikt als standaard een voorspellingsbewuste schatter (geïnspireerd door IDGI):

• In plaats van stappen te crediteren op basis van ruwe gradiënten, crediteert het elke stap op basis van de geobserveerde verandering in doelwaarschijnlijkheid ($f(h^{(s)}) - f(h^{(s-1)})$).
• Voor een scalaire waarschijnlijkheids-readout reduceert dit elk segment tot een telescopische som van eindwaarden.
• Resultaat: Volledigheid wordt gewaarborgd tot op de precisie van de drijvende-kommagetallen bij elke stapcount ($m \ge 1$). Stappen waar het model wel gradiëntgevoeligheid vertoont maar geen verandering in output geeft, worden er natuurlijk uit gefilterd, waardoor Riemann-discretisatiefouten worden geëlimineerd.

Implementatie

De methode maakt gebruik van een single-pass gebatcht algoritme dat de volledige token-per-laag kaart berekent zonder dat een backward pass (autograd) vereist is.

• Door alle interpolatiepunten in één enkele tensor te stapelen en de puntgewijze functie $f$ in één enkele forward pass te evalueren, vermijdt de methode de overhead van gradiëntberekening.
• Deze efficiëntie is specifiek mogelijk vanwege het "head-once" ontwerp, dat zorgt voor onafhankelijkheid tussen token-punten in de batch.

Belangrijkste Bijdragen

1. Exacte Additieve Decompositie: IG-Lens is de eerste methode die een exacte additieve decompositie van waarschijnlijkheid (niet logits) biedt over lagen heen, waarbij de som van de attributies gelijk is aan $\Delta p$.
2. Telescopische Identiteit: Het artikel levert een formeel bewijs (Propositie 1) dat de segmentintegralen naar de totale verandering in voorspelling telescoperen, waarbij de softmax binnen het integratiepad blijft.
3. Voorspellingsbewuste Schatting: De standaard schatter garandeert volledigheid tot op de precisie van de machine, ongeacht het aantal stappen, en filtert stappen met "gevoeligheid zonder respons" eruit.
4. Efficiënte Implementatie: Een single-pass forward evaluatie berekent de volledige attributiekaart zonder backward calls, wat de methode computationeel efficiënt maakt.
5. Positionering: De auteurs onderscheiden duidelijk waarom bestaande methoden falen in de waarschijnlijkheidsruimte (additiviteit in de logit-ruimte, bias of gebrek aan telescopie) en positioneren IG-Lens als het waarschijnlijkheidsruimte-tegenhanger van de logit-ruimte methoden.

Resultaten en Evaluatieplan

• Numerieke Verificatie: In gecontroleerde tests en op een echt model (Llama-3.2-1B-Instruct) reproduceert de standaard schatter de per-segment eindwaardeverschillen tot $\sim 10^{-7}$ (float32 precisie). De telescopische som houdt exact stand tot de floating-point precisie, onafhankelijk van het aantal stappen $m$.
• Illustratieve Output: In experimenten met de prompt "What is the capital of Vietnam?", slaagt de methode erin de attributiemassa te concentreren in de specifieke laagsegmenten waar de logit-lens waarde overgaat van laag naar hoog. De som van de attributies voor elk token komt exact overeen met $\Delta p$ tot op de floating-point precisie.
• Evaluatiestrategie: De auteurs erkennen dat "onset" (aanvang) geen intrinsieke grondwaarheid heeft. Daarom stellen zij voor om de methode te valideren via causale ablatie: als IG-Lens laag $L$ identificeert als de onset, zou het ableren van $L$ de doelwaarschijnlijkheid sterker moeten onderdrukken dan het ableren van lagen met een lagere onset-score. Deze test is ontworig om goedkoop uit te voeren en in staat om de claim te weerleggen.

Betekenis en Claims

Het artikel positioneert IG-Lens als een kwantitatieve verfijning van de iteratieve-inferentievisie. Waar de Logit Lens vraagt: "wat is de lopende voorspelling bij deze laag?", vraagt IG-Lens: "hoeveel van de uiteindelijke waarschijnlijkheid is elke laag verantwoordelijk voor?" en garandeert dat deze verantwoordelijkheden samen de hele vormt.

Bescheiden Claims en Beperkingen:

• Conditionele Marginale: De auteurs stellen expliciet dat IG-Lens een conditionele marginale meet: de extra waarschijnlijkheid die de norm+head kan uitlezen uit $h_L$ gegeven $h_{L-1}$. Het verwerpt het effect van $h_L$ dat door bovenliggende attention en MLP blokken wordt gerouteerd.
• Trade-off: De methode kiest voor "som = $\Delta p$" boven "totale causale effect via de volledige forward pass". Voor totale causale effecten blijft component-decompositie (zoals DLA) het juiste instrument.
• Niet-monotonie: De methode gaat om met niet-monotone waarschijnlijkheidstrajecten (waarbij waarschijnlijkheid wordt opgebouwd en deels weer wordt afgebroken) door gebruik te maken van gesigneerde attributies, wat een getrouwer beeld geeft van de bijdragen van lagen dan methoden gebaseerd op absolute waarden.

De centrale openstaande vraag blijft empirisch: of de IG-Lens onset de laag voorspelt waarbij causale ablatie de voorspelling het meest onderdrukt, met name voor tokens met niet-monotone trajecten.