Closed-Loop Transformers: Autoregressive Modeling as Iterative Latent Equilibrium

TL;DR

This paper introduces Equilibrium Transformers, a novel autoregressive model that iteratively refines latent representations to address the limitations of open-loop transformers, improving long-range reasoning and factual consistency.

Contribution

The paper proposes the closed-loop prediction principle and Equilibrium Transformers, integrating iterative latent refinement with theoretical guarantees and demonstrating improved performance on complex tasks.

Findings

+3.28% improvement on binary parity task

Gains up to +8.07% on difficult sequences

Unified framework for equilibrium models and diffusion language models

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.