JEPA-Reasoner: Decoupling Latent Reasoning from Token Generation

Technical Summary: JEPA-Reasoner: Decoupling Latent Reasoning from Token Generation

Problem Statement
Current Large Language Models (LLMs) operate under a unified autoregressive paradigm where high-level reasoning and low-level token generation are inextricably coupled. In this framework, the model predicts the next token based on a history of previously sampled discrete tokens. This coupling creates a fundamental bottleneck: any localized error in token selection (e.g., a grammatical slip or sub-optimal word choice) pollutes the context window, directly corrupting the trajectory of all subsequent reasoning steps. This leads to compounding errors in long-horizon tasks. Furthermore, existing attempts at latent-space reasoning (e.g., COCONUT, Recurrent Depth models) remain constrained by this coupling, often interleaving latent vectors with tokens or relying on causal masking that prevents early tokens from being informed by future reasoning states. These models also struggle to represent uncertainty, often forcing a collapse into a single reasoning path due to the discrete nature of token-level coupling.

Methodology
The authors propose JEPA-Reasoner, a novel architectural paradigm that fully decouples latent-space reasoning from token generation. The architecture consists of two distinct modules:

1. JEPA-Reasoner (The Reasoner): This module operates entirely within a continuous, normalized latent space. It utilizes a Joint-Embedding Predictive Architecture (JEPA) adapted for autoregressive latent-space generation.

   • Mechanism: It takes an input, encodes it, and uses modified Transformer blocks to predict the next latent matrix representing subsequent reasoning steps.
   • Constraint: Crucially, the Reasoner does not project outputs into vocabulary probabilities. Instead, the generated latent matrix is normalized (via a hybrid RMS and L2 normalization layer) to a unit hypersphere and looped back as input for the next step.
   • Training: It is trained via Self-Supervised Training (SST) to predict latent representations provided by a target encoder (weights updated via Exponential Moving Average). The objective is a scaled cosine similarity loss, encouraging the model to learn angular differences and consistent feature dynamics rather than discrete token prediction.

2. Talker (The Generator): This is a separate, lightweight Transformer-based module trained independently to reconstruct human-readable text.

   • Mechanism: It receives the complete latent reasoning chain generated by the JEPA-Reasoner and reconstructs the token sequence.
   • Variants: The paper introduces Mono-Talker (decoders only, for context-free tasks) and Dual-Talker (encoders and decoders, for context-aware tasks).
   • Dependency: The Talker is trained solely for reconstruction based on the Reasoner's output; it cannot make predictions independently.

Key Contributions
The paper identifies three emergent properties of this decoupled architecture:

1. Mathematical Error Containment: The authors provide a probabilistic factorization showing that because the reasoning chain ($R$) is generated independently of the token sampling ($X$), token-level errors have no mathematical pathway to influence the logic of the Reasoner ($\frac{\partial r_{t+1}}{\partial \hat{x}_t} = 0$). This isolates the reasoning engine from discrete sampling failures.
2. Continuous Guidance: The generator (Talker) has access to the entire lossless reasoning trajectory before generating any tokens, unlike coupled models where generation is interleaved with reasoning.
3. Representation of Uncertainty: The Reasoner can produce "mixed latent vectors" that encode multiple possible reasoning paths simultaneously. These vectors act as linear combinations of vocabulary embeddings, allowing the model to maintain multiple hypotheses without the overhead of beam search or multiple decodes.

Experimental Results
The authors conducted controlled experiments on synthetic and natural language tasks:

• Synthetic Tasks (Tree Search & CFG):
  • In a binary tree search task, the model achieved 99.87% accuracy. PCA analysis confirmed that predicted latent vectors formed a continuous cloud within the space spanned by discrete vocabulary embeddings, validating the "mixed latent state" hypothesis.
  • In Context-Free Grammar (CFG) generation tasks, JEPA-Reasoner demonstrated superior robustness to both input token noise (0–30% replacement) and latent space Gaussian noise compared to a standard Transformer baseline and the coupled COCONUT model.
• Natural Language Tasks (GSM8K):
  • A 0.9B parameter JEPA-Reasoner (combined with a 198M Talker) was trained on identical data and hyperparameters as a 0.9B coupled Transformer baseline.
  • Performance: The JEPA-Reasoner achieved a 51.9% accuracy on the 8-shot GSM8K benchmark, a 149.5% improvement over the baseline (20.8%).
  • Scaling: While the baseline showed a "reasoning plateau" (0.1% improvement from 5-shot to 8-shot), the JEPA-Reasoner utilized additional examples to achieve a 131.7% improvement in performance, indicating effective in-context scaling.
  • Comparison: The 0.9B model outperformed a 1B parameter CoT model (Llama 3.2) and a 3.5B parameter Recurrent Depth model (Huginn-0125) on the same benchmark.

Significance and Claims
The paper positions JEPA-Reasoner not merely as a performance boost through scaling, but as an architectural investigation into the benefits of decoupling. The authors claim that transitioning from a coupled to a decoupled paradigm yields significant gains in robustness and reasoning depth without requiring sophisticated reinforcement learning or industrial training recipes.

The work suggests that the "reasoning plateau" observed in small-scale token-based models can be overcome by moving the reasoning process into a continuous, abstract latent space. By isolating the reasoning engine from the discrete token-sampling process, the architecture provides a more robust foundation for complex reasoning, enabling error containment and the representation of uncertainty that is difficult to achieve in standard autoregressive models. The authors conclude that this decoupled approach represents a new path forward for foundation model design.