Technische Samenvatting: Tensor Memory: Vaste Grootte Recurrente Toestand voor Lang-Horizon Transformers

1. Probleemstelling

Transformers verwerken afbeeldingen en video's door ruimtelijke en temporele dimensies te platte naar lange tokenreeksen. Hoewel attentiemechanismen en KV-caching eerdere kenmerken behouden, lijden ze aan twee primaire beperkingen naarmate het reeks-horizon groeit:

1. Geheugenschaal: Praktische inferentie vertrouwt op het opslaan van eerdere sleutels en waarden, wat lineair schaalt met het aantal tokens. Dit wordt computatietechnisch en geheugenprohibitief voor lange video's en documenten.
2. Gebrek aan Persistente Toestand: Standaard Transformers missen een expliciete, persistente toestand. Ze richten zich op eerdere kenmerken, maar behouden geen gestructureerde representatie die kan worden bijgewerkt en opgevraagd naarmate de invoer verandert. Dit leidt tot kwetsbaarheid bij occlusie, ontbrekende waarnemingen, en de noodzaak om herhaaldelijk globale context te reconstrueren uit een steeds groeiende geschiedenis in plaats van een compacte, persistente toestand te behouden.

Bestaande benaderingen, zoals sparse attention, geschiedeniscompressie of retrieval-augmented contexten, verlengen vaak de contextlengte maar slaan informatie op als ongestructureerde verzamelingen van tokenvectoren of slots. Ze behouden geen expliciete ruimtelijke structuur en handhaven geen compacte "wereldtoestand" die geschikt is voor op toestand gebaseerde taken zoals video-interpretatie.

2. Methodologie: Tensor Memory

De auteurs stellen Tensor Memory voor, een lichtgewicht module die standaard Transformer-blokken aanvult met een vastgrootte recurrente 3D-geheugentensor. Deze module ontkoppelt de toestandcapaciteit van de invoerlengte terwijl een ruimtelijke inductieve bias behouden blijft.

Kernarchitectuur

Het geheugen wordt gedefinieerd als een paar volumes (verborgen toestand $h_t$ en celtoestand $c_t$) met afmetingen $C \times D \times H \times W$, waarbij de ruimtelijke dimensies constant blijven ongeacht de invoerlengte. De module werkt in een cyclus per tokenchunk:

1. Chunking: Invoertokens worden gegroepeerd in chunks. Een "leesquery" wordt afgeleid van het eerste token van de chunk, en een "schrijfsamenvatting" is een geleerde projectie van de volledige chunk.
2. Continue Adressering: Een gedeelde coördinatenpredictor (met gekoppelde gewichten) genereert continue 3D-coördinaten voor het lezen ($\mu_{read}$) en schrijven ($\mu_{write}$) binnen het voxelrooster.
3. Lezen (Continue Sampling): Het model leest uit de verborgen toestand $h_t$ bij $\mu_{read}$ met behulp van trilineaire interpolatie. Het opgehaalde kenmerk wordt geprojecteerd terug naar de tokendimensie en gefuseerd in de tokenstroom via een gegateerde residuale verbinding. Een geleerde scalair gate $\gamma$ controleert de sterkte van deze fusie, waardoor het model het geheugepad kan onderdrukken als dit niet voordelig is voor een specifieke taak.
4. Schrijven (Differentieerbare Soft Write): Het model schrijft informatie door een contentvector en een spreidingsparameter $\sigma$ te voorspellen. Het deponeert deze content als een Gaussisch-gewogen volume rond de voorspelde 3D-locatie $\mu_{write}$. Deze "soft write" vermijdt harde voxeltoewijzing, waardoor gradiënten kunnen bepalen waar informatie wordt geplaatst.
5. Recurrente Update: Het schrijfvolume wordt geconcateneerd met de huidige verborgen toestand. Een lichtgewicht, gefactoriseerde 3D-operator (dieptewijze 3D-convoluties gevolgd door een puntsgewijze mix) berekent ConvLSTM-stijl gates ($i_t, f_t, o_t, g_t$). Deze gates updaten de cel- en verborgen toestanden, waardoor informatie kan persisten, lokaal kan propageren en beschikbaar blijft zelfs wanneer waarnemingen ontbreken.

Belangrijke Ontwerpkies

• Vaste Grootte Toestand: De grootte van de geheugentensor is constant, wat de lineaire geheugengroei geassocieerd met KV-caching voorkomt.
• Ruimtelijke Structuur: Het 3D-voxelrooster behoudt expliciet ruimtelijke relaties, in tegenstelling tot ongestructureerde token-slots.
• Differentieerbare Soft Writes: Het gebruik van Gaussisch-gewogen volumes staat voor soepele, continue updates toe in plaats van discrete, harde toewijzingen.
• Gated Fusion: De geleerde gate $\gamma$ zorgt ervoor dat de module kan worden gekoppeld aan voorgeönde modellen zonder het trainen te destabiliseren; als geheugen onnodig is, onderdrukt de gate het pad.

3. Belangrijkste Bijdragen

• Nieuw Geheugenmechanisme: Introductie van een vaste grootte, recurrente 3D-voxeltoestand waarnaar tokens kunnen schrijven en van waaruit ze kunnen lezen via continue coördinaten, dienend als een expliciete ruimtelijke toestand.
• Integratie: De module integreert schoon in standaard Transformer-trainingspijplijnen en kan worden gekoppeld aan of verwijderd uit bestaande blokken zonder architecturale wijzigingen.
• Diagnostische Suite: Ontwikkeling van een gecontroleerde "toy"-diagnostische suite ontworpen om scenario's te isoleren waar persistente toestand gunstig is (bijv. coördinaatbinding, occlusie, kaartbouw) versus scenario's waar dit niet het geval is.

4. Experimentele Resultaten

De auteurs evalueerden Tensor Memory op standaard benchmarks over taal-, afbeeldings- en videodomeinen, evenals de diagnostische suite.

• Taalmodellering (WikiText-2): Tensor Memory behaalde de grootste absolute winsten, met een vermindering van de testperplexiteit van 130,30 (Base Transformer) naar 100,07. Het presteerde beter dan baselines, waaronder Transformer-XL, Local attention en Linear attention.
• Afbeelding Patch Reconstructie (CUB-200-2011): In een gemaskerde patch-reconstructietoon toonde de Tensor Memory-variant (Spatial ViT) consistente verbeteringen in PSNR en SSIM over maskerratio's (25%, 50%, 75%) in vergelijking met Base ViT, Local ViT en Registers ViT.
• Video Actieherkenning (UCF-101): Het model presteerde beter dan zowel een gezamenlijke spatiotemporale Transformer-baseline als een sterke "register token"-baseline op Top-1 (89,21% vs 82,12%) en Top-5 nauwkeurigheid. De auteurs merken op dat dit ten koste ging van aanzienlijk hogere trainingswall-clock-tijd vanwege de per-stap voxel-recurrentie.
• Toy Diagnostiek:
  • Coördinaatbinding: Tensor Memory behield een hoge nauwkeurigheid (75–91%) naarmate het aantal schrijfbewerkingen toenam, terwijl attentie-gebaseerde baselines instortten tot bijna willekeurige prestaties (7–14%).
  • No-Harm Control: Op taken die geen geheugen vereisen, bleef de geleerde gate $\sigma(\gamma)$ dicht bij zijn initialisatie (onderdrukking van het geheugepad), en behaalden alle modellen 100% nauwkeurigheid.
  • Kaartbouw: Tensor Memory presteerde aanzienlijk beter dan baselines in lang-horizon grid-navigatietaken waarbij informatie over tijd moet worden behouden.

5. Betekenis en Claims

Het artikel stelt dat Tensor Memory een praktische manier biedt om lang-horizon redenering aan te pakken door Transformers een expliciete, compacte en gestructureerde toestand te bieden.

• Ontkoppeling van Capaciteit en Lengte: Door een vaste grootte tensor te gebruiken, vermijdt het model de lineaire geheugenschaal van KV-caches, wat het theoretisch schaalbaarder maakt voor lange reeksen.
• Robuustheid tegen Occlusie: De persistente toestand stelt het model in staat een "geloofstoestand" over het tafereel te behouden, waardoor het robuuster is tegen ontbrekende waarnemingen of occlusie in vergelijking met modellen die uitsluitend vertrouwen op het reconstrueren van context uit een tokenstroom.
• Veelzijdigheid: Hoewel gemotiveerd door lang-horizon video, blijkt het mechanisme ook gunstig voor taal (accumulatie van globale context) en afbeeldingen (gestructureerde ruimtelijke aggregatie).
• Fundering voor Toekomstig Werk: De auteurs zien de huidige implementatie als een fundering. Zij suggereren toekomstige richtingen omvatten hogere dimensie geheugens (bijv. 4D voor tijd-bewuste ruimtelijke geheugen), rijkere dynamiek (bijv. optische stroom warping) en sparse updates om de efficiëntie te verbeteren.

De auteurs blijven bescheiden over de huidige beperkingen, en erkennen dat de dichte voxel-updates een doorvoerbottleneck zijn en dat de vaste grootte toestand kan verzadigen als een taak vereist het opslaan van een groot aantal verschillende details. Ze stellen echter dat het toevoegen van een kleine gestructureerde toestand een levensvatbare weg is naar capabelere lang-horizon perceptie- en redeneersystemen.