BRo-JEPA: Learning Modular Arithmetic in Latent Space

Sintesi Tecnica: BRo-JEPA – Apprendimento dell'Aritmetica Modulare nello Spazio Latente

Problematica
Il documento investiga una questione fondamentale nel deep learning: le reti neurali apprendono davvero regole algebriche astratte o si limitano a memorizzare pattern di addestramento? Mentre gli attuali modelli eccellono nei compiti matematici entro i range visti, spesso falliscono nell'estrapolare a range numerici o operazioni non viste. Gli autori sostengono che i modelli supervisionati standard e i modelli JEPA (Joint Embedding Predictive Architecture) esistenti si affidino a correlazioni statistiche piuttosto che a una logica aritmetica sottostante. Nello specifico, quando incaricati dell'aritmetica modulare (ad esempio, $y = (x + k) \mod 10$) utilizzando immagini di cifre MNIST come stati, i modelli standard si adattano alle operazioni viste ma falliscono nel generalizzare a operazioni non viste (generalizzazione zero-shot rigorosa).

Metodologia
Gli autori utilizzano il framework JEPA, in cui un modello predice la rappresentazione latente di una cifra target dato un'immagine di una cifra di contesto e un'operazione aritmetica (azione), piuttosto che predire output di pixel grezzi. Lo studio confronta diverse variazioni architettoniche:

1. Baseline Standard: Modelli supervisionati (MLP e ResNet-18) che utilizzano una perdita di cross-entropy e embedding di azione compositive additive. Questi embedding rappresentano le operazioni come multipli scalari di un vettore base (ad esempio, $op(k) = k \cdot v$) per imporre una struttura aritmetica nello spazio degli embedding.
2. JEPA con Embedding Additivi: Un modello JEPA che utilizza gli stessi embedding additivi, aumentato da una "Composition Consistency Loss". Questa perdita ausiliaria impone che l'applicazione di un'operazione composta (ad esempio, $+2$) in un unico passaggio debba produrre lo stesso risultato latente dell'applicazione sequenziale di operazioni primitive (ad esempio, $+1$ poi $+1$).
3. BRo-JEPA (Block-Rotation JEPA): Il contributo principale è un nuovo predittore che sostituisce lo standard MLP non vincolato con un predittore a rotazione a blocchi condizionato dall'operazione.
   • Meccanismo: Ispirato ai Rotary Position Embeddings (RoPE), questo predittore tratta la dinamica dello spazio latente come una serie di rotazioni planari 2D. Per un vettore latente $z$ e un'operazione $k$, la predizione è $z_p = R_{block}(k \cdot \theta) \cdot z$.
   • Bias Induttivo: Questa architettura impone esplicitamente la struttura circolare dell'aritmetica modulo-10 sullo spazio latente. Gli angoli di rotazione $\theta$ possono essere fissi (impostati a $2\pi/10$) o apprendibili.
   • Varianti: Gli autori testano la Single-Frequency Rotation (SFR), dove tutti i blocchi latenti condividono un unico angolo, e la Multi-Frequency Rotation (MFR), dove ogni blocco ha un angolo indipendente.
   • Regolarizzazione: Per prevenire il "collasso latente" (dove input diversi mappano a rappresentazioni indistinguibili), il modello impiega la regolarizzazione Variance-Invariance-Covariance (VICReg) e un target encoder tramite Exponential Moving Average (EMA).

Contributi Chiave

1. Definizione del Testbed: L'introduzione dell'aritmetica modulare su MNIST come testbed controllato per studiare l'apprendimento di regole nei modelli di mondo latenti.
2. Analisi del Fallimento: Dimostrazione che le baseline supervisionate e i modelli JEPA con embedding additivi, nonostante si adattino alle operazioni viste, falliscono nella rigorosa estrapolazione zero-shot (raggiungendo solo il 2,42% di accuratezza per ResNet-18). Essi catturano i segni delle operazioni ma falliscono nel codificare l'ampiezza o la geometria compositiva.
3. Proposta BRo-JEPA: Lo sviluppo di un modello di dinamica latente in stile JEPA caratterizzato da un predittore a rotazione a blocchi condizionato dall'operazione che si allinea con la struttura ciclica dell'aritmetica modulare.
4. Analisi Geometrica: Analisi che mostra come BRo-JEPA organizzi le rappresentazioni latenti in un manifold geometrico strutturato (un manifold circolare), consentendo al modello di interiorizzare la topologia dell'aritmetica modulare.

Risultati
Gli esperimenti sono stati condotti con addestramento sulle operazioni $\{-1, +1\}$ e valutazione sulle operazioni non viste $\{\pm 2, \dots, \pm 9\}$.

• Performance delle Baseline: Le baseline supervisionate hanno raggiunto un'alta accuratezza sulle operazioni viste (~97-98%) ma un'accuratezza quasi nulla sulle operazioni non viste (2,42% per ResNet-18). Le visualizzazioni PCA hanno rivelato che questi modelli stipano arbitrariamente le operazioni non viste nello spazio delle rappresentazioni di quelle viste.
• JEPA con Embedding Additivi: Sebbene questi modelli mostrino un'alta accuratezza di rollout (applicando ripetutamente $+1$ per simulare $+k$), falliscono nella predizione diretta zero-shot di operazioni non viste (accuratezza ~10%). La loss di consistenza compositiva ha migliorato questo dato a ~99%, ma è considerata una "zero-shot debole" poiché il modello è stato indirettamente esposto a operazioni composte durante l'addestramento.
• Performance BRo-JEPA:
  • Accuratezza Zero-Shot: Il BRo-JEPA basato su ResNet con Multi-Frequency Rotation (MFR) ha raggiunto il 99,46% di accuratezza zero-shot sulle operazioni non viste.
  • Accuratezza di Rollout: Il modello ha ottenuto anche il 99,46% di accuratezza nei rollout multi-step, confermando la stabilità delle dinamiche latenti apprese.
  • Struttura Latente: Le visualizzazioni PCA hanno confermato che BRo-JEPA riesce a costringere lo spazio latente in un manifold ellittico/circolare continuo dove le classi di cifre sono disposte sequenzialmente. La variante MFR ha prodotto prototipi di classe meglio separati (minore similarità del coseno off-diagonal) rispetto alla SFR, portando a una generalizzazione superiore.
  • Angoli Apprendibili: Quando gli angoli di rotazione erano apprendibili, la MFR ha converguto rapidamente verso una generalizzazione efficace, mentre la SFR ha mostrato ritardi significativi e instabilità, convergendo spesso a soluzioni armoniche subottimali.

Significatività e Rivendicazioni
Il documento afferma che i modelli di mondo latenti possono apprendere regole di trasformazione simbolica, ma solo quando il bias induttivo dell'architettura corrisponde alla struttura del problema. I risultati suggeriscono che:

• I predittori neurali standard mancano dei vincoli necessari per preservare la geometria compositiva per operazioni non viste.
• Modellare esplicitamente la natura ciclica dell'aritmetica modulare tramite rotazioni a blocchi nello spazio latente consente una forte generalizzazione zero-shot senza richiedere l'esposizione alle specifiche operazioni durante l'addestramento.
• Questo approccio offre una direzione promettente per l'apprendimento delle rappresentazioni oltre l'aritmetica, suggerendo che l'incorporazione di strutture di trasformazione appropriate nei modelli di transizione latente sia cruciale per l'apprendimento di regole astratte.

Gli autori rimangono modesti, osservando che i loro risultati sono specifici per l'ambiente altamente controllato dell'aritmetica modulare su MNIST con una nota struttura di gruppo ciclico. Non affermano che l'architettura generalizzerà automaticamente a compiti di ragionamento simbolico arbitrari o a gruppi di trasformazione sconosciuti senza simili priori strutturali.