BRo-JEPA: Learning Modular Arithmetic in Latent Space

Technische Samenvatting: BRo-JEPA – Leren van Modulaire Rekenkunde in de Latente Ruimte

Probleemstelling
Het artikel onderzoekt een fundamentele vraag in deep learning: leren neurale netwerken werkelijk abstracte algebraïsche regels, of memoriseren ze slechts patronen uit de trainingsdata? Hoewel huidige modellen uitblinken in wiskundige taken binnen bekende bereiken, falen ze vaak bij extrapolatie naar onbekende numerieke bereiken of operaties. De auteurs stellen dat standaard supervised modellen en bestaande Joint Embedding Predictive Architecture (JEPA) modellen vertrouwen op statistische correlaties in plaats van op onderliggende rekenlogica. Specifiek, wanneer ze worden belast met modulaire rekenkunde (bijv. $y = (x + k) \mod 10$) gebruikmakend van MNIST-cijfers als toestanden, passen standaard modellen de bekende operaties wel aan, maar falen ze in het generaliseren naar onbekende operaties (strikte zero-shot generalisatie).

Methodologie
De auteurs maken gebruik van het JEPA-framework, waarbij een model de latente representatie van een doelcijfer voorspelt op basis van een contextcijfer-afbeelding en een rekenkundige operatie (actie), in plaats van de ruwe pixeloutputs te voorspellen. De studie vergelijkt verschillende architecturale variaties:

1. Standaard Baselines: Supervised modellen (MLP en ResNet-18) die gebruikmaken van cross-entropy loss en additieve compositionele actie-embeddings. Deze embeddings representeren operaties als scalaire veelvouden van een basisvector (bijv. $op(k) = k \cdot v$) om de rekenkundige structuur in de embeddingruimte af te dwingen.
2. JEPA met Additieve Embeddings: Een JEPA-model dat dezelfde additieve embeddings gebruikt, aangevuld met een "Composition Consistency Loss". Deze extra loss dwingt af dat het toepassen van een samengestelde operatie (bijv. $+2$) in één stap hetzelfde latente resultaat oplevert als het sequentieel toepassen van primitieve operaties (bijv. $+1$ gevolgd door $+1$).
3. BRo-JEPA (Block-Rotation JEPA): De kernbijdrage is een nieuwe voorspeller die de standaard ongecontroleerde Multi-Layer Perceptron (MLP) vervangt door een operatie-geconditioneerde block-rotatie voorspeller.
   • Mechanisme: Geïnspireerd door Rotary Position Embeddings (RoPE), behandelt deze voorspeller de dynamiek van de latente ruimte als een reeks 2D-planair roterende blokken. Voor een latente vector $z$ en operatie $k$ is de voorspelling $z_p = R_{block}(k \cdot \theta) \cdot z$.
   • Inductieve Bias: Deze architectuur legt expliciet de circulaire structuur van modulo-10 rekenkunde op aan de latente ruimte. De rotatiehoeken $\theta$ kunnen vaststaan (gezet op $2\pi/10$) of leerbaar zijn.
   • Varianten: De auteurs testen Single-Frequency Rotation (SFR), waarbij alle latente blokken één hoek delen, en Multi-Frequency Rotation (MFR), waarbij elk blok een onafhankelijke hoek heeft.
   • Regularisatie: Om "latente instorting" (waarbij verschillende inputs ononderscheidbare representaties worden) te voorkomen, gebruikt het model Variance-Invariance-Covariance Regularization (VICReg) en een Exponential Moving Average (EMA) target encoder.

Belangrijkste Bijdragen

1. Definitie van Testbed: De introductie van MNIST modulaire rekenkunde als een gecontroleerd testbed voor het bestuderen van regelverwerving in latente wereldmodellen.
2. Falen-analyse: Demonstratie dat supervised baselines en JEPA-modellen met additieve embeddings, ondanks het goed fitten van bekende operaties, falen bij strikte zero-shot extrapolatie. Ze vatten de tekens van operaties wel op, maar falen in het coderen van grootte of compositionele geometrie.
3. BRo-JEPA Voorstel: De ontwikkeling van een JEPA-stijl latent dynamica model met een operatie-geconditioneerde block-rotatie voorspeller die aansluit bij de cyclische structuur van modulo-rekenkunde.
4. Geometrische Analyse: Analyse die aantoont dat BRo-JEPA de latente representaties organiseert in een gestructureerde geometrische manifold (een cirkelvormige manifold), waardoor het model de topologie van modulaire rekenkunde kan internaliseren.

Resultaten
De experimenten werden uitgevoerd met training op operaties $\{-1, +1\}$ en evaluatie op onbekende operaties $\{\pm 2, \dots, \pm 9\}$.

• Baseline Prestaties: Supervised baselines behaalden een hoge nauwkeurigheid op bekende operaties (~97-98%), maar een bijna nul-nauwkeurigheid op onbekende operaties (2,42% voor ResNet-18). PCA-visualisaties toonden aan dat deze modellen onbekende operaties willekeurig in de representatieruimte van bekende operaties propten.
• JEPA met Additieve Embeddings: Hoewel deze modellen een hoge rollout nauwkeurigheid vertoonden (het herhaaldelijk toepassen van $+1$ om $+k$ te simuleren), faalden ze bij directe zero-shot voorspelling van onbekende operaties (nauwkeurigheid ~10%). De composition consistency loss verbeterde dit naar ~99%, maar dit wordt beschouwd als "zwakke zero-shot" omdat het model tijdens de training indirect werd blootgesteld aan samengestelde operaties.
• BRo-JEPA Prestaties:
  • Zero-Shot Nauwkeurigheid: De ResNet-gebaseerde BRo-JEPA met Multi-Frequency Rotation (MFR) behaalde 99,46% zero-shot nauwkeurigheid op onbekende operaties.
  • Rollout Nauwkeurigheid: Het model behaalde ook 99,46% nauwkeurigheid in multi-step rollouts, wat de stabiliteit van de geleerde latente dynamiek bevestigt.
  • Latente Structuur: PCA-visualisaties bevestigden dat BRo-JEPA erin slaagde de latente ruimte te dwingen in een continue elliptische/cirkelvormige manifold waar cijferklassen sequentieel zijn gerangschikt. De MFR-variant produceerde beter gescheiden klasse-prototypes (lagere off-diagonal cosine similarity) vergeleken met SFR, wat leidde tot superieure generalisatie.
  • Leerbare Hoeken: Wanneer de rotatiehoeken leerbaar waren, convergeerde MFR snel naar effectieve generalisatie, terwijl SFR significante vertragingen en instabiliteit vertoonde en vaak convergeerde naar suboptimale harmonische oplossingen.

Betekenis en Claims
Het artikel claimt dat latente wereldmodellen symbolische transformatieregels kunnen leren, maar alleen wanneer de architecturale inductieve bias overeenkomt met de structuur van het probleem. De resultaten suggereren dat:

• Standaard neurale voorspellers de beperkingen missen die nodig zijn om de compositionele geometrie voor onbekende operaties te behouden.
• Het expliciet modelleren van de cyclische aard van modulaire rekenkunde via block-rotaties in de latente ruimte sterke zero-shot generalisatie mogelijk maakt zonder dat het model tijdens de training aan de specifieke operaties blootgesteld hoeft te worden.
• Deze aanpak biedt een veelbelovende richting voor representatie-leren buiten de rekenkunde, wat suggereert dat het integreren van passende transformatiestructuren in latente transitie-modellen cruciaal is voor het leren van abstracte regels.

De auteurs blijven bescheiden en merken op dat hun resultaten specifiek zijn voor de zeer gecontroleerde setting van de MNIST modulaire rekenkunde met een bekende cyclische groepstructuur. Zij beweren niet dat de architectuur automatisch zal generaliseren naar willekeurige symbolische redeneertaken of onbekende transformatiegroepen zonder vergelijkbare structurele priors.