Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

PPO for Language Models › States as Growing Contexts

Michael BrenndoerferDecember 28, 20255 min read

Language AI HandbookAll chapters

Part 37 of 67Alignment and RLHF

Chapter 279PPO for Language Models › States as Growing Contexts

Reading section 4 of 15 · States as Growing Contexts · 5 minSection 4 of 15 · 5 min

Section 4 of 15 · 5 min read

PPO for Language Models › States as Growing ContextsOne section · Sections in this chapter4 / 15 ↓

Sections in this chapter15

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

The state in language generation grows with each action. Unlike a game where the agent might return to previously visited states, or a robot navigation task where the agent can revisit locations, language generation always moves forward. Each generated token permanently extends the context. The state at step $t+1$ contains everything from state $t$ plus the newly generated token. There is no going back, no undo, and no way to visit the same state twice unless the exact same sequence of tokens has been generated before.

Think of the growing state as a conversation transcript. Once you have said something, you cannot unsay it. Every subsequent word must be consistent with everything already written. If the model commits to a particular argument in sentence one, it must carry that argument through to the end. If it introduces a character in paragraph one, that character constrains what can happen later. This irreversibility creates a basic asymmetry between early and late tokens: early choices constrain later ones, giving them outsized influence on the final response.

Formally, the state at step $t$ is:

s_t = (x_1, x_2, \ldots, x_n, y_1, y_2, \ldots, y_{t-1})

where:

$s_t$ : the state at timestep $t$ , consisting of the complete history seen so far
$x_i$ : the $i$ -th token of the prompt, for $i = 1, \ldots, n$
$n$ : the number of tokens in the prompt
$y_j$ : the $j$ -th generated token, for $j = 1, \ldots, t-1$

This formulation captures the needed nature of autoregressive generation. The state is not a compact summary of the situation; it is the complete history of what has been written. The model must condition on all of this information to decide what should come next. A word that appeared ten sentences ago might be important for maintaining coherence, while a word from three tokens ago might determine grammatical constraints on the current position. Transformers handle this beautifully through self-attention, which allows the model to selectively focus on any previous token regardless of distance.

The key insight is that this state space is effectively infinite and largely unique to each trajectory. Two different prompts lead to entirely different state spaces, and even the same prompt with different generation paths explores different states. Unlike board games where states might repeat (the same chess position can arise from different move orders) or continuous control where the system might return to similar configurations, language generation creates fresh territory with every token. Each state is a unique point in the space of possible text prefixes.

The growing state affects implementation in important ways. The model must process longer sequences as it generates tokens. Each new token requires attending to all previous tokens, which increases computation quadratically with sequence length under standard attention. This is where techniques like KV caching, which we discuss in the chapter on efficient inference, become needed for efficient generation during training rollouts. By caching the key and value projections from previous steps, the model avoids recomputing them for each new token, changing per-step complexity from quadratic to linear in the already-generated sequence.

In[6]:

Code

def demonstrate_state_growth():
    """Show how states evolve during generation."""
    prompt_tokens = [101, 2054, 2003, 1996, 3007]  # "What is the capital"

    states = []
    generated = []

    # Simulate a generation trajectory
    response_tokens = [1997, 2605, 1029, 102]  # "of France? [SEP]"

    for t, token in enumerate(response_tokens):
        current_state = prompt_tokens + generated
        states.append(
            {
                "step": t,
                "state_length": len(current_state),
                "state": current_state.copy(),
                "action": token,
            }
        )
        generated.append(token)

    return states


trajectory = demonstrate_state_growth()

Out[7]:

Console

State evolution during generation:
--------------------------------------------------
Step 0: state length = 5, action = 1997
Step 1: state length = 6, action = 2605
Step 2: state length = 7, action = 1029
Step 3: state length = 8, action = 102

Final trajectory length: 4 steps

Out[8]:

Visualization

Line chart showing how context length grows linearly with each generation step, with a shaded blue area and a red dashed boundary between prompt and generated tokens. — Monotonic state growth during the autoregressive generation process. Unlike environments with fixed state spaces or cycles, each generated token permanently extends the context. The blue shaded area shows how the total context (state) expands with each step, while the dashed red line marks where the prompt ends and generated content begins.

The generated trajectory illustrates how the context (state) expands with each step. Even a short response creates a sequence of unique states, as the growing history fundamentally changes the input to the policy at every decision point. At step 0, the policy sees only the prompt. At step 1, it sees the prompt plus one generated token. By the final step, it sees the entire conversation history. This expansion means that the policy faces a different decision problem at every timestep, even though the underlying question remains the same.