Skip to content

Preserve the sampled-token record across turns

Build a token record that keeps behavior probabilities aligned with generated actions. A small fixture separates context positions, prediction rows and loss masking before numerical comparison.

15 MIN

By InferenceInterviews · Updated

Course lessons Lesson 2 of 6

TL;DR: Keep token IDs and their positions as the record of a rollout. Tool observations remain in context even when they carry no action loss. In this fixture, generated tokens at positions 2, 3 and 6 are scored by prediction rows 1, 2 and 5.

Where you are. The workload map identifies what leaves the generator. Now specify enough evidence for another process to interpret it correctly.

Separate context from actions

Our synthetic record has seven token IDs. The first two are prompt tokens, the next two are generated, two more arrive from a tool, and the final token is generated after reading that observation. The integers are arbitrary symbols, not the output of a real tokenizer.

PositionToken IDOriginAction loss?Prediction row
0, 1101, 102PromptNoNot scored here
2, 3201, 202ModelYes1, 2
4, 5301, 302ToolNoNot scored here
6203ModelYes5

Row 5 predicts token 203 using context through position 5, including the tool's reply. Removing tool tokens would change that event. Leaving them in the input and excluding their own loss terms preserves it. The attention visibility rule and the action-loss selection rule must therefore be represented separately.

Run this standalone Python fixture. Its output gives (prediction row, target token) pairs, so the shift is visible rather than hidden inside a tensor slice.

def action_targets(ids, generated):
    if len(ids) != len(generated):
        raise ValueError("one origin flag per token")
    if any(type(flag) is not bool for flag in generated):
        raise ValueError("origin flags must be booleans")
    if generated and generated[0]:
        raise ValueError("a scored action needs a preceding context position")
    return [(i - 1, ids[i]) for i, flag in enumerate(generated) if flag]

ids = [101, 102, 201, 202, 301, 302, 203]
generated = [False, False, True, True, False, False, True]
pairs = action_targets(ids, generated)
assert pairs == [(1, 201), (2, 202), (5, 203)]
assert ids[4:6] == [301, 302]
print(pairs)

This checks indexing and ownership, not model probabilities or GPU attention. A production adapter must additionally preserve padding conventions, absolute positions and the exact prefix for each scored action.

Reconstruct one probability before aggregating

Grammar-constrained decoding provides another concrete source of allowed-token support. A syntactically valid choice still needs a probability normalized over the distribution that sampled it.

An unchanged-policy replay is a useful first check. Fix the weights and token history, apply the same sampling transformations, and compare the probability of the same observed token. Start without truncation or penalties; add each transformation separately so the first mismatch has a small search space.

ROLLOUT WORKBENCH • SYNTHETIC MODEL

Which probability did you store?

Keep the weights fixed. Change the sampler, then compare the probability of token A with the saved rollout record.

Raw model

Before sampling transforms

  1. A50.0%
  2. B30.0%
  3. C15.0%
  4. D5.0%

Temperature

Softmax at T = 1

  1. A50.0%
  2. B30.0%
  3. C15.0%
  4. D5.0%

Sampled policy

Keep top 2; normalize again

  1. A62.5%
  2. B37.5%
  3. CExcluded
  4. DExcluded
SAME WEIGHTS · TOKEN A

Recomputed probability / saved probability

0.6250 / 0.6250
1.000

The retained set held 80.0% of the temperature-adjusted mass before renormalization. Matched probabilities give a ratio of 1, as an unchanged policy should. Four illustrative tokens, fixed top-k support, no penalties or numerical noise. This is a consistency check, not a complete RL loss.

The inspector's raw distribution is [0.50, 0.30, 0.15, 0.05]. At temperature 1, retaining A and B leaves mass 0.50 + 0.30 = 0.80; A receives 0.50 / 0.80 = 0.625. Storing 0.50 while replaying 0.625 gives 0.625 / 0.50 = 1.25. The checkpoint did not move. Correctly matched probabilities give 0.625 / 0.625 = 1.

That result is a record-consistency check under the declared support, not a complete training objective. If a trainer deliberately uses a different target distribution, it needs its own explicit correction and support rules. A finite clipped ratio cannot establish that those rules were followed.

Preserve evidence through rendering and retries

Keep the original IDs even when the debugging interface shows text. Re-rendering a conversation through another template can change hidden delimiters or earlier turns. The new text may look identical while the conditioning prefix differs. Record template and tokenizer revisions with the checkpoint, and compare IDs before numerical values.

A retry also needs an identity decision. If it samples again, it is a new attempt with new probabilities; it cannot silently overwrite an older record that already received a score. If it resends an already completed result, downstream deduplication should preserve that result's identity. Keep group, rollout and turn identifiers alongside token positions so the consumer can tell the two cases apart.

Do this before moving on

Insert one extra tool token before the final action. Update the expected prediction row and explain why the earlier two action rows remain unchanged. Then deliberately mark a tool token as generated: the fixture should expose an extra training target even though the token sequence itself has not changed.

Write an acceptance check that refuses missing probabilities for generated positions, mismatched lengths and unknown tokenization revisions. Exempt context positions from the probability requirement without deleting them. Your notebook should contain the original record and the smallest changed record that your check rejects.

Go deeper

Key takeaways

  • Score each observed action against the row that predicts it.
  • Keep tool observations in context while excluding their action-loss terms.
  • Record tokenization and sampling semantics before comparing log probabilities.

Check yourself

Work it out before revealing the explanation. Include the units, assumptions and the condition that would change your answer.

  1. 1One tool token is inserted before the final action, moving that action from position 6 to 7. Which prediction row scores it?

  2. 2Tool-result tokens have no action loss. What happens to them in the model input?

Sign in to track which lessons you have finished.