Language Generation as Optimal Control: Closed-Loop Diffusion in Latent Control Space

2026-05-14

AI Summary

This work addresses the trade-off between efficiency and fidelity, irreversible error propagation, and optimization challenges inherent in autoregressive and diffusion-based language models by formulating language generation as a stochastic optimal control problem for the first time. By approximately solving the Hamilton–Jacobi–Bellman equation, the authors derive a closed-loop control policy and integrate flow matching with a global integral operator in a refined latent control space to enable efficient parallel sampling. The proposed method substantially enhances generation quality, stability, and controllability, achieving state-of-the-art performance on both unconditional language modeling and conditional generation tasks.

Abstract

This work reformulates language generation as a stochastic optimal control problem, providing a unified theoretical perspective to analyze autoregressive and diffusion models and explain their limitations (Efficiency-Fidelity Paradox, Irreversibility Error Propagation, Optimization Tractability and Fidelity) in terms of combination of trajectory singularity, adjoint state vanishing, and gradient absence. To address these issues, we approximate the solution to the Hamilton-Jacobi-Bellman (HJB) equation, yielding an optimal policy that acts as a closed-loop controller. To bypass the intractability of directly solving the HJB PDE, we employ Flow Matching as the optimal trajectory solver within the rectified latent control space. This allows our Manta-LM with Global Integral Operator to approximate the global vector field, effectively realizing a model that simultaneously achieves high-fidelity text generation and efficient, low-cost parallel sampling. Empirically, our method achieves strong performance on language modeling and conditional generation tasks, while exhibiting improved stability, efficiency, and controllability.

Problem

• language generation
• optimal control
• diffusion models
• autoregressive models
• HJB equation

Innovation

• Optimal Control
• Flow Matching
• Latent Control Space
• Hamilton-Jacobi-Bellman Equation
• Closed-Loop Generation