Technische Samenvatting: Het bestuderen van de 'Soupability' van Documenten in State Space Models

Probleemstelling

Veel real-world NLP-taken, zoals multi-document vraagbeantwoording (QA), wetenschappelijke samenvatting en juridische analyse, vereisen redeneren over volledige corpora in plaats van individuele lange documenten. Huidige Transformer-gebaseerde modellen worstelen hiermee vanwege de prohibitieve $O(L^2)$ attention-kosten en de inflexibiliteit van monolithische encoding-strategieën. In standaardbenaderingen worden alle documenten geconcateneerd tot één enkele sequentie voordat ze worden verwerkt; consequent vereist elke wijziging aan het corpus (bijv. het toevoegen of verwijderen van een enkel document) het opnieuw coderen van de gehele input vanaf nul. Dit verhindert de modulaire hergebruik van documentrepresentaties en creëert een bottleneck voor schaalbare, dynamische retrieval-systemen.

Hoewel Structured State Space Models (SSM's) zoals Mamba en Mamba2 lineaire tijd-sequentieverwerking bieden en inputs comprimeren in fixed-length verborgen toestanden, worden ze doorgaans ingezet met dezelfde monolithische concatenatie-strategie, waardoor ze dezelfde beperkingen qua hercoderen erven.

Methodologie: Document Souping

De auteurs stellen Document Souping voor, een strategie geïnspireerd door model souping (het middelen van gewichten van gefinetunede checkpoints), maar toegepast op verborgen toestanden. De kernhypothese is dat de lineaire recursie-mechanismen in SSM's de onafhankelijk gecodeerde documentrepresentaties post hoc mogelijk maken om te worden samengevoegd via eenvoudige lineaire operaties (bijv. middeling) om een samengestelde toestand te vormen die geschikt is voor downstream redeneren.

De Architectuur

1. Onafhankelijke Encoding: Elk document $d_i$ in een corpus wordt onafhankelijk gecodeerd door een gedeelde SSM-encoder ($SSM_\theta$), wat een set laag-specifieke verborgen toestanden $\{h^{(l)}_i\}$ produceert.
2. Pooling: Deze per-document toestanden worden op elke laag geaggregeerd met behulp van een commutatieve pooling-operator (bijv. element-wise gemiddelde, som of max) om een enkele "souped" toestand $h^{(l)}_{soup}$ te creëren.
3. Decoding: De gepoolde toestanden $\{h^{(1)}_{soup}, \dots, h^{(L)}_{soup}\}$ worden samen met de query $q$ in de decoder geïnjecteerd om het antwoord $\hat{y}$ te genereren.

Training en Optimalisatie

• Finetuning: Het artikel demonstreert dat gepre-trainde SSM's "souping" niet van nature ondersteunen; volledige encoder-decoder finetuning is vereist om het model te leren de gepoolde toestanden effectief te interpreteren.
• Geheugenefficiëntie: Om de geheugenoverhead van het backpropageren door meerdere onafhankelijke document-encoders te beheersen, gebruiken de auteurs document-niveau activation checkpointing. Dit herberekent de forward-activaties tijdens de backward pass, waardoor het geheugengebruik constant blijft, ongeacht het aantal documenten ($k$).
• Pooling Strategie: Uit uitgebreide experimenten blijkt dat element-wise middeling zonder normalisatie de meest robuuste en stabiele aggregatiemethode is, die superieur presteert aan sommatie (dat lijdt onder onbegrensde magnitude-groei) en max-pooling.

Belangrijkste Bijdragen

1. Systematische Empirische Studie: De eerste uitgebreide investigatie naar "soupability"—het vermogen van SSM's om onafhankelijk gecodeerde documentrepresentaties te mergen voor multi-document redeneren.
2. Prestatievalidatie: Demonstratie dat gefinetunede Mamba2-modellen met souped representaties competitieve of superieure prestaties leveren op multi-hop QA (HotpotQA), long-document QA (RACE, QuALITY) en sparse retrieval (NIAH) vergeleken met standaard concatenatie-baselines.
3. Schaalbaarheid en Generalisatie: Bewijs dat de aanpak schaalt naar honderden documenten (getest tot 256) en generaliseert naar variërende contextgroottes, mits het model is gefinetuned op een voldoende aantal segmenten.
4. Architecturale Specificiteit: Bewijs dat deze modulaire aanpak uniek geschikt is voor SSM's. Een analoge "cache souping" experiment op Transformers (LLaMA 3) faalde significant, wat benadrukt dat de fixed-size, recurrente toestand van SSM's een cruciale enabler is voor deze techniek.
5. Inference Efficiëntie: Een nieuwe inference workflow waarbij documenten één keer worden gecodeerd, gecached en dynamisch worden samengesteld, wat substantiële latency-besparingen biedt (tot 14,6x versnelling voor 32k tokens) vergeleken met het opnieuw coderen van geconcateneerde contexten.

Experimentele Resultaten

De auteurs evalueerden Mamba2-8B en Mamba2-2.7B modellen over verschillende benchmarks:

• Multi-Hop QA (HotpotQA): Full encoder-decoder gefinetunede modellen met average pooling bereikten 47.8 EM op 5 documenten, waarmee ze decoder-only finetuning overtroffen en de concatenatie-baselines evenaarden of overtroffen. De methode bleef robuust naarmate het aantal distractors toenam.
• Long-Document QA (RACE/QuALITY): Souping presteerde aanzienlijk beter dan traditionele concatenatie-benaderingen op long-document benchmarks, vooral wanneer de trainingsconfiguratie overeenkwam met de test-tijd soup grootte.
• Sparse Retrieval (NIAH): Modellen getraind op grotere aantallen segmenten (bijv. 8 segmenten) behielden een hoge nauwkeurigheid (71.45% EM) bij tests op 16 segmenten, terwijl modellen getraind op minder segmenten faalden in generalisatie.
• Schaalbaarheid: Hoewel een model getraind op 64 documenten faalde in extrapolatie naar 256 documenten, maakte voortgezette finetuning op 128 documenten het model in staat om 37.25 EM op 256 documenten te bereiken, wat bewijst dat de architectuur zich kan aanpassen aan enorme contexten.
• Latency: Voor een 32k token context duurde het verwerken van een query met gecachte souped toestanden 65.8 ms, vergeleken met 963.2 ms voor de standaard concatenatie methode.

Betekenis en Claims

Het artikel claimt dat document souping een nieuw paradigma ontsluit voor modulaire, schaalbare corpus-redenering. Door de document-encoding te ontkoppelen van de query-verwerking, maakt de aanpak het volgende mogelijk:

• Efficiënte Updates: Documenten kunnen één keer worden gecodeerd en gecached, wat dynamische corpus-updates en retrieval mogelijk maakt zonder de volledige input opnieuw te verwerken.
• Schaalbare RAG: De methode dient als een drop-in vervanging voor de stap "concateneer de opgehaalde passages" in Retrieval-Augmented Generation (RAG) systemen, waarbij een variabel maar begrensd aantal documenten efficiënt wordt ondersteund.
• Architecturale Fit: Dit werk stelt vast dat de lineaire recursie-eigenschappen van SSM's hen uniek in staat maken om state-level compositie te ondersteunen, een capaciteit die standaard Transformers missen vanwege hun afhankelijkheid van een variabele-lengte KV-cache.

De auteurs concluderen dat hoewel extrapolatie naar ongeziene soup-groottes uitdagend kan zijn, passende finetuning-strategieën SSM's in staat stellen om effectief grote, dynamische documentensets te hanteren, wat een lichtgewicht en effectieve oplossing biedt voor long-context redeneren.