ViT-Linearizer: Distilling Quadratic Knowledge into Linear-Time Vision Models

Dit paper introduceert ViT-Linearizer, een kennisdistillatieframework dat de kwadratische complexiteit van Vision Transformers overbrengt naar efficiënte, lineaire recurrente modellen, waardoor aanzienlijke snelheidswinsten worden behaald en de prestaties van Mamba-architecturen op visuele taken worden verbeterd.

De kern

Stel je voor dat je een zeer slimme, maar traag werkende meesterkunstenaar hebt: de ViT (Vision Transformer). Deze kunstenaar kan elke foto in detail bekijken, alsof hij elke pixel met elke andere pixel vergelijkt. Hij is ongelooflijk goed in het begrijpen van beelden, maar er zit een groot nadeel aan: hij werkt kwadratisch.

Wat betekent dat?

• Als je een kleine foto van 100 pixels geeft, moet hij 10.000 vergelijkingen maken.
• Als je een hoge-resolutie foto van 10.000 pixels geeft, moet hij 100.000.000 vergelijkingen maken!
• De computer wordt hierdoor snel moe, traag en heeft veel geheugen nodig. Het is alsof je een gigantische bibliotheek moet doorzoeken voor elk klein detail.

Aan de andere kant hebben we de RNN/Mamba-modellen. Dit zijn slimme, maar snellere "snelheidslopers". Ze kijken naar een foto alsof ze een verhaal lezen: woord voor woord, pixel voor pixel. Ze zijn lineair (ze worden niet veel trager als de foto groter wordt), maar ze missen soms de diepe, globale connecties die de meesterkunstenaar wel heeft. Ze zijn snel, maar niet altijd even slim.

De oplossing: ViT-Linearizer
De auteurs van dit paper hebben een slimme truc bedacht: kennis distilleren. Ze willen de wijsheid van de trage, super-slimme meester (ViT) overdragen op de snelle, maar minder ervaren leerling (Mamba), zodat de leerling net zo slim wordt, maar wel blijft rennen.

Hoe doen ze dit? Met twee creatieve methoden:

1. "Kijk eens hoe ik kijk" (Activatie Matching)

Stel je voor dat de meesterkunstenaar (ViT) naar een foto van een hond kijkt. Hij denkt: "Oh, die vlek hier is een oor, en die hier is een staart, en ze horen bij elkaar." Hij maakt een mentale kaart van hoe alles met elkaar samenhangt.

De leerling (Mamba) kijkt ook naar de hond, maar hij heeft die globale kaart niet. Hij ziet alleen losse stukjes.
De ViT-Linearizer dwingt de leerling om niet alleen naar het eindresultaat te kijken, maar ook naar de tussentijdse gedachten van de meester.

• De analogie: Het is alsof de meester zijn notitieboekje openlegt en zegt: "Kijk, bij dit stukje van de foto, moet je ook denken aan dat andere stukje. Kijk hoe mijn hersenen die twee verbinden."
• De leerling moet zijn eigen "mentale kaart" (activaties) laten lijken op die van de meester. Hierdoor leert de snelle leerling plotseling hoe hij verbanden moet leggen, zonder dat hij zelf die trage, zware berekeningen hoeft te doen.

2. "Het raadsel van de gemaskerde delen" (Masked Prediction)

Nu komt het tweede deel. Stel je voor dat je de meester vraagt: "Wat zie je op deze foto?" en hij antwoordt direct. Dat is saai voor de leerling; hij leert niet echt.

In plaats daarvan bedekken ze een groot deel van de foto met een deken (maskeren).

• De opdracht: De meester kijkt naar de hele foto en zegt wat er onder de deken zit. De leerling kijkt alleen naar de zichtbare delen en moet raden wat er onder de deken zit, gebaseerd op wat de meester zou zeggen.
• De analogie: Het is alsof je een puzzel oplost. Als je alleen maar het antwoord ziet, leer je niets. Maar als je moet raden wat er ontbreekt, moet je echt nadenken over de context.
• Door deze "raadsels" op te lossen, wordt de leerling veel slimmer in het begrijpen van de context van een beeld. Hij leert niet alleen wat hij ziet, maar ook wat er misschien is.

Het resultaat

Door deze twee methoden te combineren, krijgen we een model dat:

1. Net zo slim is als de trage meester (hij haalt bijna dezelfde hoge scores op tests).
2. Veel sneller is (tot wel 4 keer sneller bij grote, hoge-resolutie foto's).
3. Minder geheugen nodig heeft, waardoor het op gewone computers of zelfs telefoons kan werken.

Kortom:
De auteurs hebben een manier gevonden om de "grote brein" van een trage supercomputer over te dragen naar een "snel brein". Het is alsof je de wijsheid van een oude, wijsgeerige professor overbrengt op een jonge, energieke student. De student blijft snel en actief, maar denkt nu net zo diep na als de professor. Dit maakt het mogelijk om complexe beeldtaken (zoals het analyseren van medische scans of zelfrijdende auto's) veel efficiënter uit te voeren.

Technische samenvatting

Probleemstelling

Vision Transformers (ViTs) hebben revolutionaire prestaties geleverd in visuele taken dankzij hun vermogen om globale zelf-attention (self-attention) te gebruiken. Deze mechanismen zorgen voor robuuste token-gebaseerde representaties, zelfs in zwak gesuperviseerde scenario's. Echter, de belangrijkste beperking van ViTs is hun kwadratische rekencomplexiteit ($O(L^2)$) ten opzichte van de sequentielengte $L$. Dit wordt een kritieke bottleneck bij het verwerken van hoogwaardige, hoogresolutie beelden of lange contexten, wat leidt tot onaanvaardbare inferentietijden en hoge hardware-eisen.

Hoewel alternatieven zoals recurrente token-mixers (bijv. Mamba, RWKV, xLSTM) lineaire complexiteit ($O(L)$) bieden en efficiënter zijn, missen deze modellen vaak de rijke representatieve capaciteiten van ViTs. Tot nu toe zijn recurrente visiemodellen beperkt gebleven tot kleinere datasets en modelgroottes, en directe kennisoverdracht (distillatie) van ViTs naar deze lineaire modellen resulteerde vaak in zwakke studentmodellen.

Methodologie: ViT-Linearizer

De auteurs introduceren ViT-Linearizer, een cross-architectuur distillatieframework dat de kennis van een kwadratische ViT-leraar (teacher) overbrengt naar een lineaire, recurrente student (bijv. Adventurer met Mamba-2). De methode bestaat uit twee kerncomponenten:

1. Activeringsmatching (Activation Matching):

   • In plaats van alleen de einduitvoer te vergelijken, focust deze methode op de intermediate activation maps van de ViT. Deze maps reflecteren de token-tot-token afhankelijkheden die door de zelf-attention zijn geleerd.
   • De methode berekent voor zowel de leraar als de student de paarsgewijze cosinus-ähnlichkeiten tussen alle tokens in meerdere tussenliggende lagen.
   • Een $\ell_2$-verliesfunctie minimaliseert de afstand tussen de genormaliseerde activeringskaarten van de leraar en de student. Dit dwingt de recurrente student om de precieze lokale representaties en afhankelijkheden van de ViT na te bootsen, ondanks het lineaire rekenmodel.

2. Gemaskerde Predictie (Masked Prediction):

   • Voor de eindlaag worden de features niet direct gedistilleerd. In plaats daarvan wordt een gemaskerde strategie toegepast (geïnspireerd door MAE).
   • De student ontvangt een ingevoerde afbeelding waarbij een deel van de patch-tokens is vervangen door een leerbare [mask]-token, terwijl de leraar de volledige afbeelding ontvangt.
   • De student moet de representaties van de onzichtbare (gemaskerde) tokens voorspellen door de output van de [mask]-tokens te alignen met de output van de leraar op die posities.
   • Belangrijke nuance: Activeringsmatching wordt alleen toegepast op de zichtbare tokens om informatielekken te voorkomen die zouden ontstaan als de student de antwoorden voor gemaskerde tokens zou "kijken" tijdens het matchen.

Het totale verlies is een combinatie van activeringsmatching en gemaskerde predictie: $L = L_{act} + \lambda L_{mask}$.

Belangrijkste Bijdragen

• Kennisoverdracht van Kwantiteit naar Kwaliteit: Het paper toont aan dat het mogelijk is om de "kwadratische kennis" (token-afhankelijkheden) van ViTs effectief over te dragen naar lineaire recurrente modellen zonder significante prestatieverlies.
• Nieuwe State-of-the-Art voor Recurrente Modellen: Door ViT-Linearizer te gebruiken, verbetert de prestatie van Mamba-gebaseerde architecturen (Adventurer) aanzienlijk, zelfs op standaard benchmarks zoals ImageNet.
• Efficiëntie bij Hoogresolutie: De methode lost het hardware-probleem van ViTs op bij lange sequenties (hoogresolutie beelden) door de inferentiecomplexiteit lineair te houden, terwijl de nauwkeurigheid behouden blijft.

Resultaten

De auteurs evalueren hun methode op diverse taken en datasets:

• ImageNet Classificatie:

  • Een gedistilleerd Adventurer-Base model (Mamba-2) bereikte 84,3% top-1 nauwkeurigheid op ImageNet-1k. Dit is een verbetering ten opzichte van eerdere supervised Mamba-modellen (82,6%) en zelfs beter dan sommige ViT-baselines.
  • Bij een grotere input (448x448) bereikte het model 85,0% nauwkeurigheid, wat een nieuwe state-of-the-art is voor deze familie van recurrente modellen.
  • Snelheidswinst: Er werd een 2,1x versnelling in inferentietijd bereikt vergeleken met de ViT-leraar op ImageNet, zonder noemenswaardige nauwkeurigheidsverlies.

• Semantische Segmentatie (Hoge Resolutie):

  • Op ADE20K (512x512) bereikte het model een mIoU van 51,3%, een verbetering ten opzichte van de ViT-leraar (51,0%) en andere baselines, met een 2,74x snelheidswinst.
  • Op Cityscapes (512x1024, zeer lange sequenties) was de snelheidswinst nog indrukwekkender: 4,21x snellere inferentie vergeleken met de ViT-leraar, terwijl de nauwkeurigheid zelfs steeg (82,0% vs 80,2% voor de baseline).

• Kwalitatieve Analyse:

  • Visuele inspectie van activatiekaarten toont aan dat de gedistilleerde student veel minder "ruis" heeft dan een standaard gesuperviseerd recurrent model. De patronen lijken sterk op die van de ViT-leraar, wat aantoont dat de student de vermogens van de zelf-attention heeft overgenomen.

Betekenis en Impact

ViT-Linearizer biedt een brug tussen de theoretische efficiëntie van lineaire recurrente modellen en de praktische effectiviteit van grote Transformers.

• Schaalbaarheid: De voordelen van de methode nemen toe naarmate de sequentielengte (resolutie) toeneemt, wat het ideaal maakt voor toekomstige toepassingen in hoogresolutie beeldverwerking en lange context-taken.
• Paradigmaverschuiving: Het paper suggereert een nieuwe transfer-learning paradigma waarbij complexe, dure modellen worden gebruikt voor pre-training, waarna de kennis wordt overgebracht naar efficiëntere, lineaire modellen voor inferentie. Dit maakt het mogelijk om de voordelen van foundation models te benutten op hardware met beperkte middelen.
• Toekomstperspectief: De resultaten tonen aan dat recurrente visiemodellen, wanneer correct gedistilleerd, concurren kunnen met ViTs, wat de weg vrijmaakt voor bredere adoptie van lineaire architecturen in visuele AI.