Answer-token-only loss¶
One supervised token per example¶
The readout looks at one position. Training puts the loss on that position and nowhere
else. build_example holds this invariant:
input_ids = encode(full_prompt(state, criterion)) + [answer_token]
labels = [-100] * len(encode(full_prompt(...))) + [answer_token]
answer_token = answer_slots(tokenizer, prompt, n_options)[0][decision.answer_index]
-100 is the index PyTorch's cross-entropy ignores, so every prompt token is masked
and the only gradient comes from predicting the answer letter after Answer:.
Teaching the model to reproduce ticket text, or to write an explanation it will never be asked for, spends a small model's capacity on the wrong thing. The answer token is the one produced by answer-slot discovery, so training supervises exactly the token the readout will read.
build_example only needs an object with an encode method, so the invariant is
tested in CI with a stub tokenizer and no torch:
from typedecide import Criterion, Decision, Option
from typedecide.training import IGNORE_INDEX, build_example
class CharTokenizer:
"""One token per character. Enough to see the masking."""
def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
return [ord(c) for c in text]
criterion = Criterion("queue", "Which queue?", (
Option("billing", "Billing"), Option("shipping", "Shipping"),
))
decision = Decision(id="T-1", state="Charged twice.", criterion=criterion,
answer_id="shipping")
example = build_example(CharTokenizer(), decision, max_length=1024)
supervised = [t for t in example["labels"] if t != IGNORE_INDEX]
print(supervised, chr(supervised[0])) # [66] B
Truncation comes from the left of the state¶
When encode(prompt) plus the answer token exceeds max_length, something has to go.
Ordinary right-truncation would cut off the criterion, the options and Answer:. That
destroys the answer slot: the supervised position would land on a token that is not an
answer letter.
So build_example removes tokens from the left of the state:
- The tail (criterion, options,
Answer:) is protected and never touched. - The instruction preamble (
INSTRUCTIONplusEvidence:) is kept at the front when there is room. - The oldest part of the state is dropped, keeping the evidence nearest the criterion.
Edge cases, all handled explicitly:
| Situation | Behaviour |
|---|---|
The tail alone needs more than max_length - 1 tokens |
TrainingError: raise max_length or shorten the option descriptions |
| No room for the preamble and any evidence | Warning logged; only the state tokens nearest the criterion are kept |
| The tokenizer does not encode the head as a token-level prefix of the prompt | Warning logged; the whole prompt is truncated from the left, which still preserves the tail |
Evaluation never truncates
TorchReadout scores over-length prompts in full and logs a warning when a
max_length was given, because truncating at scoring time would move or destroy
the answer boundary. A state that was truncated in training is seen whole at
evaluation. If that mismatch matters to you, shorten or summarise states before
loading.
validate warns about long states using an estimate
The long_state finding uses a tokenizer-free estimate (the larger of about 1.3
tokens per word and one token per four characters) against a 1024-token budget,
so validation never downloads a model. It is a smoke alarm. Its message says the
end of the ticket is what gets lost; for training in this library it is the
start of the state that is dropped, as described above.
Padding: right for training, left for scoring¶
The two are different on purpose.
- Training right-pads (
pad_batch). Real tokens keep positions0..n-1, which is what default position ids give, and padding carries-100labels and a zero attention mask. - Scoring left-pads and passes explicit
position_ids = (attention_mask.cumsum(-1) - 1).clamp(min=0). The readout takes the logits at index-1; with right padding that index is a pad token for every row except the longest. And without explicit position ids, a left-padded row is told its first real token sits at positiongap, so a row's answer would change depending on which other prompts shared its batch.
Both mistakes produce wrong numbers silently, not a crash, which is why they are spelled out in the source.