Thinking Inside The Box The Implicit Chain Transformer For Efficient State Tracking

Cerebras Skip to main content

Cerebras Announces First Quarter 2026 Results >>

Dec 12 2025 Thinking Inside the Box: The Implicit Chain Transformer for Efficient State Tracking Jayadeep More Arun Kumar Patil Ganesh Venkatesh

Motivation Large Language Model (LLM) decoders have demonstrated remarkable capabilities in open-ended generation, reasoning, and human-computer interaction. However, the standard autoregressive formulation suffers from a  representational bottleneck : to generate the next token, the model must implicitly re-derive the underlying semantic context by attending to the entire history. This  statelessness  renders standard Transformers surprisingly brittle on tasks necessitating the maintenance of a running state—such as calculating the sum of a list of numbers modulo X or performing graph traversal. In this work, we introduce the  Implicit Chain Transformer (ICT) , a novel architecture designed to bridge this gap. By propagating a learnable "intent" latent vector forward across time steps, our method enables the model to explicitly  update and contextualize  a running state, rather than solely relying on the re-derivation of attention patterns over the history. Early evaluations demonstrate that the ICT achieves strong accuracy on state-intensive toy tasks - identified as challenging for Transformers by prior work - without incurring the inference latency costs  associated with Chain-of-Thought (CoT) prompting. This work represents a foundational step towards our broader objective: enabling Transformers to perform efficient, robust reasoning entirely within latent space. The Challenge: The Stateless Bottleneck Standard Transformers suffer from a fundamental  representational bottleneck : they are stateless. To generate token t, the model must re-derive the current context by attending to all previous tokens. It cannot simply "remember" the current state (e.g., "sum is 5"); it must implicitly recalculate it at every step. This makes them brittle on tasks requiring deep sequential dependency, leading to hallucinations as context grows. Task Setup: Modulo Arithmetic & Graph Traversal We targeted two tasks where "approximate" attention can face challenges as noted by prior work [1]: Sum Modulo X:  Requires maintaining a precise running total. A single "carry" error corrupts the entire future sequence. In theory, this is a parallel reduction task where each token can independently compute the answer instead of building on the previous partial outcome. However, as the sequence length increases, this form of re-computation at every token becomes too expensive for the model. Graph Traversal:  Requires tracking a path through a network. The model must respect connectivity, creating strict sequential dependence across tokens. This task places a much higher bar on the quality of state tracking and propagation, given the tight reliance across time steps.

The Proposal: Implicit Chain Transformer We introduce the  Implicit Chain Transformer (ICT) . Unlike standard models that restrict information flow to vertical (layer-to-layer) movement, the ICT propagates a learnable  "Intent Vector" (zt)  horizontally across time steps. This vector acts as a compressed working memory: the model "writes" the current logical state into zt​ and passes it to t+1, maintaining continuity without re-processing the entire history. We investigate two distinct strategies for this state propagation: 1. Autoregressive Intent Propagation (Dense)  In this formulation, the state is updated continuously. For every single token generation step, we capture the latent vector from the  final  layer of token t and inject it into the  early  layers of token t+1. This mirrors standard autoregressive decoding but adds a persistent "memory stream" alongside the token stream, allowing the model to micro-manage state at the word level. 2. Periodic Intent Propagation (Sparse)  Here, we decouple state updates from token generation. We inject special  tokens at regular intervals in the input stream. Mechanism:  We restrict intent propagation to occur  only  at these  boundaries. Intuition:  This forces the model to treat these tokens as "semantic checkpoints", aggregating the preceding context into a coherent summary before moving on. By focusing state updates only where the model is explicitly developing high-level concepts, we align the intent propagation architecture with the logical structure of the data rather than the arbitrary tokens. This approach also offers significant advantages for inference optimization, which we will detail in a future post.

Training with "Iterative Latent Feedback"  The primary challenge with recurrence is that it typically breaks the parallel training efficiency that makes Transformers scalable. To enable intent propagation without reverting to the slow, sequential training of RNNs, we introduce a  multi-pass approximation : Parallel Pass:  We process the entire sequence in parallel (standard Transformer mode) to generate initial latent representations. Feedback Injection:  We take the latent vector from the  final  layer of the first pass, project it via an MLP, and fuse it as the "intent" input for the  early  layers. Refinement Pass:  We perform a second forward pass with this fused context to compute the final loss.

While this introduces a fixed computational overhead (requires additional forward passes), the cost is  constant regardless of sequence length. This preserves the O(1) sequential complexity of Transformer training, avoiding the crippling O(N) training time associated with Backpropagation Through Time (BPTT) in RNNs. Early Evaluation We evaluate on the above tasks using a standard GPT-2 style decoder-only model. For these evaluations, we use following network configuration: Implicit Chain Transformer Variations: We layer our modifications on the baseline model as follows: 1. Intent Vector : A new, dedicated MLP network called the GenMLP reads the output from an intermediate(last-2nd layer for 8 layer decoder model ) or final layer(3-decoder-layer model) of the LLM . 2. Intent Propagation : The intent Vector is fused with the output of the first decoder Layer using a dedicated Fuse-Intent network 3. GenMLP network : This network processes the intermediate output to generate the Intent Vector . Architecture is as below num_hidden_layers=2 expansion_factor=4 activation=gelu norm_type=layernorm use_residual=True

4. Fuse-Intentnetwork : This Network achieves the...