State-to-Go Transformers

• State-to-Go Transformers are transformer-based models that make the sequential state explicit through mechanisms like stack-buffer states, latent transitions, or compressed key-value caches.
• They are applied in diverse domains such as transition-based parsing, visual reinforcement learning, and bounded-state decoding to enhance performance and interpretability.
• Exposing task-relevant state improves practical outcomes by clarifying history versus pending information, optimizing efficiency and inference speed across applications.

Searching arXiv for the provided works and closely related terminology to ground the article. State-to-Go Transformers are transformer-based sequence models in which the evolving state of a process is made explicit rather than left entirely implicit in unrestricted attention. In the cited literature, this designation appears in three closely related forms: dynamic stack-and-buffer-aware decoding for transition-based parsing, an offline pretrained State-to-Go (STG) Transformer for learning from visual observation, and a bounded-state decoder obtained by interpreting a transformer as a multi-state RNN with a compressed key-value cache (Astudillo et al., 2020, Zhou et al., 2023, Oren et al., 2024). Taken together, these formulations suggest a broader family of methods in which the model conditions on what has been committed, what remains pending, or what state must be carried forward across time.

1. Conceptual scope and recurrent motif

The three usages of the term differ in task, architecture, and training objective, but all introduce an explicit state variable into transformer computation.

Formulation	State representation	Core mechanism
Transition-based parsing	Parser stack $S_t$ and buffer $B_t$	Specialized cross-attention heads with dynamic masks and optional position embeddings
LfVO STG Transformer	Latent visual transition $(e_t,e_{t+1})$ and discriminator score	Offline adversarial transition modeling plus intrinsic reward generation
Bounded-state decoder	Per-layer key/value cache $S_t=(K_{1:t},V_{1:t})$	Cache compression policy, especially TOVA, to keep fixed-size state

In transition-based parsing, the relevant state is symbolic and structured: some subset of the input tokens lives on the stack and the remainder is in the buffer. In visual reinforcement learning, the state is latent and predictive: embeddings of observed frames are used to model and discriminate transitions. In long-context decoding, the state is the accumulated key-value cache itself, interpreted as a recurrent hidden state with potentially unbounded size (Astudillo et al., 2020, Zhou et al., 2023, Oren et al., 2024).

This comparison suggests that “state-to-go” is not a single canonical transformer architecture. Rather, it denotes a design pattern in which transformer inference is reorganized around an explicit sequential state that can be masked, predicted, scored, or compressed.

2. Transition-based parsing as explicit stack-buffer state modeling

In transition-based parsing, modeling the parser state is central to performance. Astudillo et al. modify a vanilla sequence-to-sequence Transformer so that decoding is conditioned on the current parser configuration. If the input sentence is $w=(w_1,\dots,w_N)$ and the encoder produces $H=[h_1;\dots;h_N]\in\mathbb R^{N\times d}$, then a standard cross-attention head uses

$Q=b_tW^Q,\quad K=HW^K,\quad V=HW^V$

and

$$e_{t,i}=\frac{Q\cdot K_i^T}{\sqrt d},\qquad \alpha_{t,i}=\operatorname{softmax}_i(e_{t,\cdot}),\qquad \text{head}_t=\sum_{i=1}^N \alpha_{t,i}V_i.$$

The stack-Transformer alters this mechanism by reserving two decoder heads: one attends only to the stack $S_t$, and the other only to the buffer $B_t$. For the stack head,

$B_t$0

$B_t$1

$B_t$2

with $B_t$3 if $B_t$4, and $B_t$5 otherwise. A corresponding buffer head uses its own mask, position embeddings, and projection matrices. All remaining heads are free to attend to the full encoder representation. By remasking and reindexing positional embeddings at each step, the method mirrors the stack-LSTM’s ability to push and pop in $B_t$6 time (Astudillo et al., 2020).

The same paper also introduces a more local state signal through a SHIFT-label multi-task objective. When the parser emits SHIFT at time $B_t$7, it also predicts one of the $B_t$8 most frequent words or tags of the token being shifted; that label becomes part of the action history seen by the decoder self-attention. This mechanism is described as akin to using contextual Bi-LSTM features for only the top of the stack or buffer, and it helps the decoder “know” what just entered the stack (Astudillo et al., 2020).

Empirically, the explicit state modeling improves both dependency and AMR parsing. On dev-set averages $B_t$9, the vanilla Transformer obtains PTB $(e_t,e_{t+1})$0 LAS/UAS and AMR2.0 $(e_t,e_{t+1})$1 Smatch; adding SHIFT-100 multi-task yields PTB $(e_t,e_{t+1})$2 and AMR2.0 $(e_t,e_{t+1})$3; the stack-Transformer reaches PTB $(e_t,e_{t+1})$4 and AMR2.0 $(e_t,e_{t+1})$5. On low-data AMR1.0, the stack-Transformer shows $(e_t,e_{t+1})$6 Smatch versus vanilla and $(e_t,e_{t+1})$7 versus the multi-task variant. With model size 2-layer, the stack-Transformer incurs only $(e_t,e_{t+1})$8 drop relative to 12-layer, while vanilla loses $(e_t,e_{t+1})$9 points. On the test set, the best variant reports PTB UAS/LAS $S_t=(K_{1:t},V_{1:t})$0, AMR2.0 Smatch $S_t=(K_{1:t},V_{1:t})$1, and AMR1.0 Smatch $S_t=(K_{1:t},V_{1:t})$2; the paper characterizes these as competitive with SOTA $S_t=(K_{1:t},V_{1:t})$3, second best published, and best published without extra recategorization, respectively (Astudillo et al., 2020).

3. Offline pretrained STG Transformer for learning from visual observation

The State-to-Go Transformer of Ren et al. addresses learning from visual observation, where policies must be recovered from only visual observation data. The framework has two stages. In Stage 1, an STG Transformer is pretrained offline to predict and differentiate latent transitions of demonstrations. In Stage 2, the pretrained model provides intrinsic rewards for downstream reinforcement learning, and the agent learns merely from intrinsic rewards (Zhou et al., 2023).

The model operates on visual “states” $S_t=(K_{1:t},V_{1:t})$4, each defined as a short stack of raw frames, such as four $S_t=(K_{1:t},V_{1:t})$5 grey-scale Atari frames or a single $S_t=(K_{1:t},V_{1:t})$6 RGB Minecraft frame. A convolutional feature encoder $S_t=(K_{1:t},V_{1:t})$7 with $S_t=(K_{1:t},V_{1:t})$8 maps each $S_t=(K_{1:t},V_{1:t})$9 to a latent token embedding

$w=(w_1,\dots,w_N)$0

A length-$w=(w_1,\dots,w_N)$1 sequence $w=(w_1,\dots,w_N)$2 is combined with learnable positional embeddings and fed into a small GPT-style causal transformer $w=(w_1,\dots,w_N)$3 with three stacked multi-head causal self-attention layers, 4 heads, and hidden size $w=(w_1,\dots,w_N)$4. The transformer predicts a residual update $w=(w_1,\dots,w_N)$5, producing

$w=(w_1,\dots,w_N)$6

In parallel, the method trains a 1-Lipschitz discriminator $w=(w_1,\dots,w_N)$7, implemented as an MLP with spectral normalization and weight-clipping, and a temporal-distance regressor $w=(w_1,\dots,w_N)$8 that predicts the signed symlog delay

$w=(w_1,\dots,w_N)$9

The TDR head is used only in pretraining as an auxiliary loss (Zhou et al., 2023).

Pretraining uses an expert-only dataset $H=[h_1;\dots;h_N]\in\mathbb R^{N\times d}$0 of visual trajectories. The discriminator is optimized with a Wasserstein-GAN adversarial loss

$H=[h_1;\dots;h_N]\in\mathbb R^{N\times d}$1

with 1-Lipschitzness enforced by weight-clipping $H=[h_1;\dots;h_N]\in\mathbb R^{N\times d}$2. The encoder and transformer minimize a generator loss consisting of an adversarial term and an $H=[h_1;\dots;h_N]\in\mathbb R^{N\times d}$3 matching term,

$H=[h_1;\dots;h_N]\in\mathbb R^{N\times d}$4

and the TDR objective is

$H=[h_1;\dots;h_N]\in\mathbb R^{N\times d}$5

The total pretraining objective is

$H=[h_1;\dots;h_N]\in\mathbb R^{N\times d}$6

Positive pairs are expert transitions $H=[h_1;\dots;h_N]\in\mathbb R^{N\times d}$7 and negative “fake” pairs are $H=[h_1;\dots;h_N]\in\mathbb R^{N\times d}$8; because the discriminator is trained purely offline, no online negative samples are required during pretraining (Zhou et al., 2023).

After pretraining, the intrinsic reward for online reinforcement learning is the discriminator’s WGAN gap:

$H=[h_1;\dots;h_N]\in\mathbb R^{N\times d}$9

The closer the real transition is to expert behavior relative to the generator’s prediction, the larger the intrinsic reward. Stage 2 freezes $Q=b_tW^Q,\quad K=HW^K,\quad V=HW^V$0, $Q=b_tW^Q,\quad K=HW^K,\quad V=HW^V$1, and $Q=b_tW^Q,\quad K=HW^K,\quad V=HW^V$2, initializes a policy $Q=b_tW^Q,\quad K=HW^K,\quad V=HW^V$3, and runs PPO, or PPO+SIL in Minecraft, using only the intrinsic reward (Zhou et al., 2023).

The empirical evaluation covers Atari and Minecraft. Atari tasks are Breakout, Freeway, Qbert, and Space Invaders; the metric is mean episodic return over 100 evaluation episodes averaged over 4 seeds. Offline STG hyperparameters are learning rate $Q=b_tW^Q,\quad K=HW^K,\quad V=HW^V$4, batch 16, MSE weight 0.5, Adv weight 0.3, TDR weight 0.1, 3 CSL layers, 4 heads, embedding dim 512. Mean final return is reported as follows: Breakout, GAIfO 1.5, ELE 22.0, STG 288.8, Expert 212.5, PPO 274.8; Freeway, GAIfO 0.6, ELE 2.7, STG 21.8, Expert 31.9, PPO 32.5; Qbert, GAIfO 394.4, ELE 4698.6, STG 27234.1, Expert 15620.7, PPO 14293.3; SpaceInvaders, GAIfO 260.2, ELE 384.6, STG 502.1, Expert 1093.9, PPO 942.5. In Minecraft, on “Pick a flower,” “Milk a cow,” “Harvest tallgrass,” and “Gather wool,” STG outperforms ELE by 50–400% across all four tasks; the example given is STG $Q=b_tW^Q,\quad K=HW^K,\quad V=HW^V$5 versus ELE $Q=b_tW^Q,\quad K=HW^K,\quad V=HW^V$6 in “milk a cow” (Zhou et al., 2023).

4. Decoder-only transformers as bounded multi-state RNNs

A separate but related line of work conceptualizes decoder-only transformers as multi-state RNNs. In this formulation, the hidden state at token step $Q=b_tW^Q,\quad K=HW^K,\quad V=HW^V$7 is the concatenated key/value cache

$Q=b_tW^Q,\quad K=HW^K,\quad V=HW^V$8

where $Q=b_tW^Q,\quad K=HW^K,\quad V=HW^V$9 and $$e_{t,i}=\frac{Q\cdot K_i^T}{\sqrt d},\qquad \alpha_{t,i}=\operatorname{softmax}_i(e_{t,\cdot}),\qquad \text{head}_t=\sum_{i=1}^N \alpha_{t,i}V_i.$$0. At time $$e_{t,i}=\frac{Q\cdot K_i^T}{\sqrt d},\qquad \alpha_{t,i}=\operatorname{softmax}_i(e_{t,\cdot}),\qquad \text{head}_t=\sum_{i=1}^N \alpha_{t,i}V_i.$$1, the model computes

$$e_{t,i}=\frac{Q\cdot K_i^T}{\sqrt d},\qquad \alpha_{t,i}=\operatorname{softmax}_i(e_{t,\cdot}),\qquad \text{head}_t=\sum_{i=1}^N \alpha_{t,i}V_i.$$2

updates the multi-state by appending the new key/value,

$$e_{t,i}=\frac{Q\cdot K_i^T}{\sqrt d},\qquad \alpha_{t,i}=\operatorname{softmax}_i(e_{t,\cdot}),\qquad \text{head}_t=\sum_{i=1}^N \alpha_{t,i}V_i.$$3

and forms the context vector

$$e_{t,i}=\frac{Q\cdot K_i^T}{\sqrt d},\qquad \alpha_{t,i}=\operatorname{softmax}_i(e_{t,\cdot}),\qquad \text{head}_t=\sum_{i=1}^N \alpha_{t,i}V_i.$$4

This yields a recurrence

$$e_{t,i}=\frac{Q\cdot K_i^T}{\sqrt d},\qquad \alpha_{t,i}=\operatorname{softmax}_i(e_{t,\cdot}),\qquad \text{head}_t=\sum_{i=1}^N \alpha_{t,i}V_i.$$5

which the paper describes as exactly an RNN-style recurrence except that the state is a matrix of every past key/value (Oren et al., 2024).

To obtain a bounded multi-state RNN, the cache size is fixed by setting $$e_{t,i}=\frac{Q\cdot K_i^T}{\sqrt d},\qquad \alpha_{t,i}=\operatorname{softmax}_i(e_{t,\cdot}),\qquad \text{head}_t=\sum_{i=1}^N \alpha_{t,i}V_i.$$6 and applying a compression policy $$e_{t,i}=\frac{Q\cdot K_i^T}{\sqrt d},\qquad \alpha_{t,i}=\operatorname{softmax}_i(e_{t,\cdot}),\qquad \text{head}_t=\sum_{i=1}^N \alpha_{t,i}V_i.$$7 whenever the temporary cache reaches size $$e_{t,i}=\frac{Q\cdot K_i^T}{\sqrt d},\qquad \alpha_{t,i}=\operatorname{softmax}_i(e_{t,\cdot}),\qquad \text{head}_t=\sum_{i=1}^N \alpha_{t,i}V_i.$$8. The paper lists Window, Window$$e_{t,i}=\frac{Q\cdot K_i^T}{\sqrt d},\qquad \alpha_{t,i}=\operatorname{softmax}_i(e_{t,\cdot}),\qquad \text{head}_t=\sum_{i=1}^N \alpha_{t,i}V_i.$$9, Headscores, and TOVA. TOVA, short for Token Omission Via Attention, computes the current attention weights

$S_t$0

and omits the single token $S_t$1 with the smallest $S_t$2. The layer-wise pseudocode then removes $S_t$3, restores state size $S_t$4, computes

$S_t$5

and outputs $S_t$6 (Oren et al., 2024).

The same work sketches an approximation-error bound. If

$S_t$7

then

$S_t$8

If $S_t$9 is the lowest dropped weight, then

$B_t$0

This gives a direct control term in terms of the discarded attention mass (Oren et al., 2024).

The bounded-state interpretation yields a practical State-to-Go Transformer for streaming or segmented inference. The procedure initializes an empty state for each layer, processes each segment token by token, carries forward only the bounded per-layer state of size $B_t$1, and stores only the $B_t$2 key/value pairs per layer between segments. Implementation pointers include pre-allocating per-layer buffers of size $B_t$3 for $B_t$4 and $B_t$5, shifting out one row and appending one row in $B_t$6, and implementing the TOVA “find-min” from the most recent attention vector $B_t$7 (Oren et al., 2024).

Experimentally, TOVA is reported to outperform several baseline compression policies on four long-range tasks and several LLMs. Results are nearly on par with the full model while using in some cases only $B_t$8 of the original cache size, which translates to 4.8X higher throughput. The detailed results state that on PG-19, TOVA is within 0.5 perplexity of full context using $B_t$9 to $B_t$00 of cache; on SQuALITY, within 1 ROUGE-point at $B_t$01–1024; on QASPER, within 1 F1-point at $B_t$02; and on story generation, a loss of 19% at $B_t$03 and 6% at $B_t$04. The paper also reports up to 88% reduction in KV-cache memory, characterized as an eight-fold batch-size increase, and states that for any given $B_t$05, TOVA dominates Window, Window$B_t$06, and Headscores (Oren et al., 2024).

5. Shared principles across parsing, RL, and long-context inference

Across these formulations, explicit state enters transformer computation in three distinct but structurally analogous ways. In transition-based parsing, the state is partitioned into “what has been committed” and “what remains to be generated,” realized as stack and buffer masks over encoder states. In visual observation learning, the state is a latent transition manifold in which expert and predicted transitions are separated by a discriminator and temporally aligned by a TDR auxiliary loss. In bounded-state decoding, the state is the key-value cache itself, treated as a recurrent hidden state that can be compressed without retraining (Astudillo et al., 2020, Zhou et al., 2023, Oren et al., 2024).

This suggests a common design principle: transformer models can be strengthened by exposing task-relevant state variables that would otherwise be implicit in unrestricted self-attention or cross-attention. In the parsing setting, the effect is to recover global stack-LSTM-style and local Bi-LSTM-style state signals while remaining fully differentiable and parallelizable. In the RL setting, the state representation is learned so that discriminator scores correlate tightly with expert-like transitions, and those scores become dense intrinsic rewards. In the bounded-decoder setting, the state is made explicit enough that one can reason about recurrence, hidden-state size, and approximation error (Astudillo et al., 2020, Zhou et al., 2023, Oren et al., 2024).

A further implication is that “state-to-go” methods sit at an interface between transformer sequence modeling and recurrent state-space reasoning. The parsing paper explicitly connects its design to stack-LSTMs and contextual Bi-LSTMs; the multi-state RNN paper states that decoder-only transformers can be conceptualized as unbounded multi-state RNNs; and the RL paper uses a transformer to learn transition dynamics and a discriminator to score progress in latent state space. Although the mechanisms are different, all three formulations impose structure on the sequential state rather than treating it as an undifferentiated history (Astudillo et al., 2020, Zhou et al., 2023, Oren et al., 2024).

6. Extensions, limitations, and unresolved questions

The parsing formulation is presented as readily extensible beyond the reported dependency and AMR experiments. The paper states that the same pattern could be generalized to structured-generation tasks in which the decoder must consult a dynamic memory of items pending versus processed. Examples given are AMR with richer graph-substructure states such as separate heads for pending reentrancies, transition-based generation of SQL queries or code with heads for pending arguments or scopes, and machine translation with heads for translated versus untranslated source. It also raises the possibility of adding a recursive composition sub-network so that stack elements carry composed subgraph embeddings, thereby bringing back the best of stack-LSTMs together with transformer parallelism and scaling advantages (Astudillo et al., 2020).

The visual-observation STG Transformer is accompanied by explicit limitations. The work trains a separate STG for each Atari game, although it states that a larger “multi-task” STG also works nearly as well. Learning a single universal STG across hundreds of games remains open. Scalability to very high-resolution video, long horizons, or 3D first-person view at full resolution would require larger transformers and more compute. Hierarchical “multi-step” STG rewards at different temporal scales are proposed as a future direction, and the paper notes that pretraining dataset bias can limit exploration beyond the demonstrated behavior if expert videos are suboptimal or biased. The ablation reported for removing TDR shows both slower convergence and lower final returns, which the authors use to motivate the auxiliary temporal-alignment objective (Zhou et al., 2023).

The bounded-state decoder formulation exposes a different limitation regime: a direct speed-memory-quality trade-off. The paper states that smaller $B_t$07 implies faster inference and less memory but eventually degrades accuracy. TOVA improves this frontier, but it does not eliminate the underlying constraint that fixed-state compression must discard information once $B_t$08. The work nonetheless presents the bounded-state view as useful both conceptually, by relating transformers to RNN-style recurrence, and practically, by mitigating the key-value cache bottleneck in long-context and streaming inference (Oren et al., 2024).

Taken together, these limitations indicate that State-to-Go Transformers are best understood as a family of structured transformer modifications whose effectiveness depends on how well the chosen state variable matches the task. In parsing the state is externally defined by the transition system; in LfVO it is learned from expert-only visual trajectories; and in long-context decoding it is inherited from the transformer’s own cache dynamics. This suggests that the central open problem is not only how to model state efficiently, but also how to choose the state abstraction that preserves the task-relevant information needed for action selection, generation, or inference.