Le rappresentazioni anisotrope migliorano la pianificazione nei modelli del mondo JEPA

Abstract degli autori

I modelli latenti del mondo apprendono dinamiche condizionate dalle azioni nello spazio delle rappresentazioni e spesso valutano le azioni candidate in base alla distanza euclidea dalla rappresentazione di un obiettivo. L’addestramento congiunto di norma regolarizza la rappresentazione per evitarne il collasso, ma la geometria che ne risulta determina anche il peso degli errori terminali durante la pianificazione. Mostriamo che predizioni accurate e rappresentazioni non collassate non garantiscono un costo di pianificazione nello spazio latente allineato al compito: la regolarizzazione gaussiana isotropa può indurre una geometria che ordina gli esiti realizzabili in modo diverso rispetto al costo del compito. Per affrontare questa discrepanza, introduciamo AnisoWM con $Λ$Reg, che sostituisce il target gaussiano isotropo fisso con una matrice di covarianza diagonale apprendibile, soggetta a vincoli di traccia fissa e di anisotropia. L’obiettivo di predizione, l’architettura del predittore e il pianificatore euclideo restano invariati; il target viene usato solo durante l’addestramento. La nostra analisi descrive come la predizione determina l’allocazione della varianza del target, come tale allocazione dipende dalla distribuzione dei dati di addestramento e a quali condizioni la metrica indotta riduce il regret della pianificazione. In quattro ambienti di controllo visivo, AnisoWM migliora il successo della pianificazione rispetto a LeWorldModel in tutti e quattro. Anche il suo costo di pianificazione nello spazio latente mostra una maggiore concordanza con gli esiti del compito. Sito del progetto: https://rkdrn79.github.io/AnisoWM-page/