Latent Reasoning in Continuous Space for Unified Multimodal Models

Technical Breakdown§

Unified Multimodal Models (UMMs) like BAGEL combine vision-language understanding and image generation but struggle to leverage chain-of-thought (CoT) reasoning for image generation due to a modality mismatch: reasoning occurs in discrete token space while image generation operates in continuous space. LARC bridges this gap by performing reasoning over continuous hidden states.

Core Methodology§

LARC introduces a latent reasoning framework with two training stages:

1. Curriculum Supervised Fine-Tuning (SFT): Gradually transitions from token-level reasoning to latent reasoning. The model is first trained on standard CoT data (discrete tokens) and then progressively replaces token predictions with continuous latent states. This is achieved by feeding the hidden state at each reasoning step as input to the next step, with a curriculum that increases the proportion of latent steps over time.

1. Information-Gain Reinforcement Learning (RL): Self-evolves latent reasoning without ground truth supervision. The reward is based on information gain—the reduction in uncertainty about the target image given the latent reasoning. Formally, the model maximizes the mutual information between latent reasoning states $z$ and the generated image $x$:

$$\mathcal{L}_{RL} = \mathbb{E}_{z \sim \pi_\theta}[\log p(x|z) - \log p(x)]$$

where $p(x|z)$ is the likelihood of the image given the latent reasoning, and $p(x)$ is the marginal likelihood. This encourages latent states that are informative about the image.

Implementation Details§

• The UMM backbone is a transformer that processes both text and image tokens. For latent reasoning, the model outputs a sequence of continuous hidden states $z_1, ..., z_T$ instead of discrete tokens.
• During curriculum SFT, the model is trained with a mix of discrete and continuous reasoning steps. The loss is a weighted combination of token-level cross-entropy and latent-state reconstruction loss.
• In RL stage, the policy is the latent reasoning process, and the reward is computed as the information gain. The RL objective is optimized using policy gradient methods (e.g., PPO) with a baseline to reduce variance.
• Test-time scaling: LARC supports parallel (multiple latent reasoning paths) and sequential (longer reasoning chains) scaling, demonstrating improved performance with more compute.

Code Snippet§

import torch
import torch.nn as nn

class LARCReasoner(nn.Module):
    def __init__(self, base_model, latent_dim):
        super().__init__()
        self.base = base_model  # UMM backbone
        self.latent_proj = nn.Linear(base_model.hidden_size, latent_dim)
        self.to_token = nn.Linear(latent_dim, base_model.vocab_size)  # for SFT

    def forward(self, input_ids, images, mode='latent'):
        # mode: 'discrete' for token reasoning, 'latent' for continuous
        hidden = self.base.encode(input_ids, images)
        if mode == 'discrete':
            logits = self.base.lm_head(hidden)
            return logits
        else:
            # latent reasoning: iterate over continuous states
            z = hidden[:, -1]  # initial state
            latents = []
            for _ in range(self.num_reasoning_steps):
                z = self.latent_proj(z)
                latents.append(z)
                # feed z back as input (e.g., via cross-attention)
                hidden = self.base.decode_step(z, hidden)
            return latents

Results§

• LARC improves generation quality, prompt following, and compositional fidelity across benchmarks.
• Achieves state-of-the-art on GenEval with a 7% point gain over BAGEL.
• Demonstrates test-time scaling: performance improves with more parallel or sequential reasoning steps, unlike token-based generation.

Conclusion§

LARC effectively bridges the modality gap by reasoning in continuous space, enabling UMMs to inherit the test-time scaling and controllability of LLMs for image generation.