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

TL;DR

This paper introduces a novel framework for language generation using optimal control theory, specifically employing closed-loop diffusion in latent space to improve efficiency, fidelity, and controllability.

Contribution

It reformulates language generation as a stochastic optimal control problem and proposes a practical solution using Flow Matching within a latent control space.

Findings

Achieves high-fidelity text generation with efficient, low-cost sampling.

Demonstrates improved stability and controllability in language modeling.

Performs strongly on 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.