Technische Samenvatting: Interdomain Attention

Probleemstelling

Transformers en diepe State Space Models (SSM's) vertegenwoordigen twee tegenovergestelde ontwerpparadigma's voor sequentiemodeling. Standaard attentiemechanismen (Vaswani et al., 2017) maken gebruik van een groeiende key-value (KV) cache, waarbij queries worden gerouteerd via inhoudsgebaseerde matching. Hoewel dit een scherpe recall biedt, leidt dit tot kwadratische rekenkosten ($O(N_q N)$) en lineaire geheugengroei ($O(N)$) met betrekking tot de sequentielengte. Daarentegen comprimeren diepe SSM's (zoals S4, Mamba) de volledige context tot een recurrente toestand met een vaste grootte, die in $O(1)$ tijd per stap wordt bijgewerkt. Deze compressie gaat echter ten koste van directe query-key matching; de vaste toestand is niet direct adresseerbaar door specifieke queries, wat vaak leidt tot zwakkere prestaties in taken die een precieze ophaalactie of inhoudsgebaseerde attention vereisen.

Hybride architecturen hebben geprobeerd deze benaderingen te combineren, en lineaire attention-varianten hebben gezocht naar een directe vermindering van de attention-kosten. Dit artikel stelt een andere route voor: het integreren van een SSM direct in het attentiemechanisme om de efficiëntie van een toestand met vaste grootte te behouden, terwijl elke query de gecomprimeerde geschiedenis kan raadplegen.

Methodologie: Interdomain Attention

De kernbijdrage is Interdomain Attention, een tokenmixer die kernel-attention en state space models verenigt. De methode is afgeleid van een kernel-regressievisie op attention en een HiPPO-stijl basisprojectie.

Theoretische Afleiding

1. Kernel Feature Map: De attention-kernel $K(q, k)$ wordt benaderd door een eindige feature map $\xi(q)^\top \xi(k)$.
2. Basisprojectie: In plaats van alle keys en values op te slaan, projecteert de methode de key-features $\xi(k)$ en values $v$ op een gedeelde set basisfuncties die wordt onderhouden door een enkele SSM-recurrentie (specifiek HiPPO-stijl). Dit resulteert in een toestand met vaste grootte die onafhankelijk is van de sequentielengte.
3. Gecomprimeerde Coëfficiënten: De integralen van deze projecties leveren drie toestandscomponenten op:
   • $U$: Basisprojectie-coëfficiënten voor keys (inducerende variabelen).
   • $\Gamma$: Value-projecties.
   • $\eta$: Normaliserende constanten.
4. Query-geconditioneerde Readout: Op elke positie wordt de query $q$ afgebeeld op de feature-ruimte $\xi(q)$. De output wordt berekend door te attenderen op de gecomprimeerde coëfficiënten in plaats van op individuele tokens:
   $$\tilde{o} = \frac{\sum_m (\xi(q)^\top u_m) \gamma_m}{\sum_m (\xi(q)^\top u_m) \eta_m}$$
   In matrixvorm: $\tilde{O} = \frac{F_q U^\top \Gamma}{F_q U^\top \eta}$, waarbij $F_q$ de query-featurematrix is.

Schaalbare Implementatie

De praktische implementatie is een geleerde relaxatie van de theoretische afleiding, ontworpen voor stabiliteit en efficiëntie:

• Feature Map: Gebruikt een DeltaNet-stijl neurale feature map $\xi(x) = \text{SiLU}(\text{SC}(x)) / \|\text{SiLU}(\text{SC}(x))\|^2$, waarbij SC een causale depthwise convolutie is.
• Recurrente Kern: Benadert de HiPPO-projectie met behulp van een complex-diagonale S4D-recurrentie. De SSM-input concateneert de key-feature $\xi(k)$ en de value $v$.
• Normalisatie: Past input RMSNorm toe met leerbare per-head biases op de key- en value-helften voordat deze de SSM binnenkomen.
• Readout zonder Noemer: Om het mengen van incompatibele schalen veroorzaakt door input-rescaling te vermijden, wordt de normaliserende term $\eta$ weggelaten, wat resulteert in een lineaire attention-readout: $\tilde{O} = F_q U^\top \Gamma$.
• Architectuur: Geïntegreerd in een Llama-stijl decoder met pre-normalisatie, SwiGLU feedforwards en rotary position embeddings (RoPE).

Belangrijkste Bijdragen

1. Constructie: Een query-geconditioneerde tokenmixer met vaste toestand, afgeleid van kernel-regressie en HiPPO-basisprojectie, met een duidelijk onderscheid tussen de ideale theoretische afleiding en de schaalbare implementatie.
2. Mechanisme Decompositie: Een ablatiestudie bij 125M parameters identificeert de query-geconditioneerde projectie als de dominante factor die prestatiewinst boven standaard S4D-mixers drijft. De dubbele key/value-input draagt bij aan een kleinere, secundaire winst.
3. Schaalstudie: Een uitgebreide evaluatie op FineWeb-Edu van 125M tot 1,3B parameters onder een gelijk recurrente-toestand-budget.

Experimentele Resultaten

De studie vergelijkt Interdomain Attention met een uitsluitend S4D-tokenmixer en een standaard softmax-attention-baseline.

• Prestaties vs. S4D: Interdomain Attention verbetert de S4D-tokenmixer op elke schaal (125M–1,3B), met een validatieperplexiteit die 13–16% lager is.
• Prestaties vs. Softmax: Bij 1,3B parameters overtreft Interdomain Attention de softmax-baseline met hetzelfde recept op validatieperplexiteit (7,98 vs. 8,63) en op een acht-taken suite voor gezond verstand (+3,03 procentpunten gemiddelde nauwkeurigheid). Bij 125M is het ongeveer gelijk aan softmax.
• Lengte-extrapolatie: Door gebruik te maken van de kern met vaste toestand, vertoont Interdomain Attention "lengte-vlak" gedrag. Terwijl de perplexiteit van de softmax-baseline significant verslechtert buiten de trainingscontext (4K), behoudt Interdomain stabiele prestaties tot 14K (3,5× de trainingscontext).
• Recall-beperkingen: Het artikel merkt op dat Interdomain Attention, net als andere modellen met vaste toestand, worstelt met exacte string-associatieve recall (bijvoorbeeld naald-in-hooiberg-taken) in vergelijking met softmax-attention. Dit wordt geïdentificeerd als een structurele beperking van compressie met vaste toestand en niet als een specifiek gebrek aan het Interdomain-mechanisme.

Betekenis en Claims

Het artikel beweert dat Interdomain Attention succesvol de kloof overbrugt tussen de efficiëntie van SSM's met vaste toestand en de inhoudsgebaseerde ophaalcapaciteiten van attention. Door features te projecteren op een SSM-basis, stelt het queries in staat de gecomprimeerde geschiedenis te raadplegen zonder de kwadratische kosten of lineaire geheugengroei van standaard attention.

De primaire betekenis ligt in het aantonen dat een query-geconditioneerde projectie op een vaste recurrente toestand zowel pure SSM's als, op grotere schaal, standaard softmax-attention kan overtreffen, terwijl het de gewenste eigenschappen van lengte-extrapolatie van SSM's erft. De auteurs suggereren dat deze benadering een haalbare weg biedt voor taalmodellering met lange context, waarbij geheugenefficiëntie en consistente prestaties over variërende contextlengtes cruciaal zijn. Toekomstige werkrichtingen die worden genoemd, zijn het integreren van sterkere kernen met vaste toestand (bijvoorbeeld Mamba-3) of het combineren van de readout met snelle weight-update-regels.