Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

JEPA-Reasoner: Decoupling Latent Reasoning from Token Generation

Bingyang Kelvin Liu, Ziyu Patrick Chen, David P. Woodruff

TL;DR

JEPA-Reasoner introduces a decoupled architecture that performs latent-space reasoning independent of token generation. A Talker module translates latent trajectories into tokens, enabling pure latent reasoning with robust autoregressive generation. The approach yields strong gains on synthetic tasks and real-world GSM8K results, demonstrating reduced error propagation and the potential for multi-threaded latent reasoning. Overall, the work provides a principled framework for separating high-level reasoning from surface-token generation with practical implications for scalable reasoning in language tasks.

Abstract

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.

JEPA-Reasoner: Decoupling Latent Reasoning from Token Generation

TL;DR

JEPA-Reasoner introduces a decoupled architecture that performs latent-space reasoning independent of token generation. A Talker module translates latent trajectories into tokens, enabling pure latent reasoning with robust autoregressive generation. The approach yields strong gains on synthetic tasks and real-world GSM8K results, demonstrating reduced error propagation and the potential for multi-threaded latent reasoning. Overall, the work provides a principled framework for separating high-level reasoning from surface-token generation with practical implications for scalable reasoning in language tasks.

Abstract

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.
Paper Structure (55 sections, 4 equations, 7 figures, 8 tables)

This paper contains 55 sections, 4 equations, 7 figures, 8 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Joint-Embedding Predictive Architectures (JEPA).
  4. Latent Space Reasoning.
  5. Autoregressive Models and Robustness.
  6. Model Architecture
  7. JEPA-Reasoner
  8. Model Components.
  9. Latent Space Generation.
  10. Training Objective and Target Encoder.
  11. Action-Taker Model
  12. Training Procedure
  13. Pretraining
  14. Objective and Methodology.
  15. Self-Supervised Training
  16. ...and 40 more sections

Figures (7)

  • Figure 1: Architecture of JEPA-Reasoner and its action taker (Talker). The Reasoner consists of an embedding layer as the token encoder and Transformer blocks as the predictor. The embedding layer for input tokens always uses the latest weights, while the weight of the embedding layer for target tokens is the exponential moving average of the historical weights of the input embedding layer.
  • Figure 2: The relative performance of JEPA-Reasoner $R$ and Transformer baseline $T$ at each generation step with different fractions of correct tokens in the input sequence being replaced with wrong ones.
  • Figure 3: Visualization of the tree structure in the given example.
  • Figure 4: A picture demonstrating how CFG sequence is generated. It involves replacing non-terminal symbols at each level with symbols from the next level according to a given rule.
  • Figure 5: CFG production rule used to generate training and test samples in \ref{['sec:robust-reasoner']}. This rule gives sequence lengths ranging from about 600 symbols to 700 symbols.
  • ...and 2 more figures