JEPA-Reasoner: Decoupling Latent Reasoning from Token Generation
This work addresses a problem in AI for improving reasoning and generation in models, but it appears incremental as it builds on existing JEPA and latent reasoning approaches.
The paper tackles the limitations of Joint-Embedding Predictive Architecture (JEPA) lacking generative abilities and Transformer models suffering from compounding errors in token-by-token generation by proposing JEPA-Reasoner, which decouples latent reasoning from token generation to achieve superior robustness to compounding error.
While Joint-Embedding Predictive Architecture (JEPA) has emerged as a powerful architecture for learning rich latent representations, it fundamentally lacks generative abilities. Meanwhile, latent space reasoning attempts for Transformer models like COCONUT do improve performance, but they ultimately rely on token-by-token generation, which still accumulates compounding error and relies on context information to gain reasoning insights. To address these limitations, we propose JEPA-Reasoner, a novel JEPA model enhanced with generative ability that reasons in latent space. We augment it with a separate action-taker model, Talker, to produce human-readable sentences. Our approach demonstrates that decoupling latent space reasoning and token generation enables JEPA-Reasoner to produce mixed latent vectors that might lay the foundation for multi-threaded reasoning, while performing autoregressive generation with superior robustness to compounding error.