Closed-Loop Transformers: Autoregressive Modeling as Iterative Latent Equilibrium
Akbar Anbar Jafari, Gholamreza Anbarjafari
TL;DR
This work identifies an inherent open-loop bottleneck in autoregressive transformers and proposes Equilibrium Transformers (EqT), which perform iterative latent refinement to reach a self-consistent equilibrium before emitting each token. By introducing a learned energy function that enforces predictive coherence, bidirectional reconstruction, and memory grounding, EqT achieves approximate MAP inference with convergence guarantees and improved performance on hard long-range tasks. The approach unifies deep equilibrium models, diffusion language models, and test-time training under a single closed-loop framework, and initial parity-task experiments show noteworthy gains where one-shot inference struggles. The results suggest that inference-time deliberation can qualitatively enhance reasoning, grounding, and long-context fidelity, offering a principled path toward more robust language models with controllable compute.
Abstract
Contemporary autoregressive transformers operate in open loop: each hidden state is computed in a single forward pass and never revised, causing errors to propagate uncorrected through the sequence. We identify this open-loop bottleneck as a fundamental architectural limitation underlying well-documented failures in long-range reasoning, factual consistency, and multi-step planning. To address this limitation, we introduce the closed-loop prediction principle, which requires that models iteratively refine latent representations until reaching a self-consistent equilibrium before committing to each token. We instantiate this principle as Equilibrium Transformers (EqT), which augment standard transformer layers with an Equilibrium Refinement Module that minimizes a learned energy function via gradient descent in latent space. The energy function enforces bidirectional prediction consistency, episodic memory coherence, and output confidence, all computed without external supervision. Theoretically, we prove that EqT performs approximate MAP inference in a latent energy-based model, establish linear convergence guarantees, and show that refinement improves predictions precisely on hard instances where one-shot inference is suboptimal. The framework unifies deep equilibrium models, diffusion language models, and test-time training as special cases. Preliminary experiments on the binary parity task demonstrate +3.28% average improvement on challenging sequences, with gains reaching +8.07% where standard transformers approach random performance, validating that the benefit of deliberation scales with task difficulty. Just as attention mechanisms resolved the sequential bottleneck of recurrent networks, we propose that closed-loop equilibrium may resolve the commitment bottleneck of open-loop autoregression, representing a foundational step toward language models.