Teacher forcing — next-token targets

Background

A decoder-only LM is trained by shifting the sequence: the input is every token except the last, and the target is every token except the first. Each position predicts the next token, so a length-L sequence yields L-1 supervised predictions — a dense signal on (almost) every token.

$$\text{input} = x_{0:L-1}, \qquad \text{target} = x_{1:L}$$

Problem statement

Implement causal_lm_targets(sequences) returning (inputs, targets, num_predictions) for a batch.

Input

• sequences — array-like of shape (B, L) of integer token IDs (L ≥ 2).

Output

Returns (inputs, targets, num_predictions):

• inputs = sequences[:, :-1] — shape (B, L-1),
• targets = sequences[:, 1:] — shape (B, L-1),
• num_predictions = B * (L-1).

Examples

Example 1 — one sequence

Input:  sequences = [[1, 2, 3, 4]]
Output: inputs = [[1, 2, 3]], targets = [[2, 3, 4]], num_predictions = 3

Example 2 — batch of two

Input:  sequences = [[1, 2, 3], [4, 5, 6]]
Output: inputs = [[1, 2], [4, 5]], targets = [[2, 3], [5, 6]], num_predictions = 4

Constraints

• inputs drops the last column; targets drops the first.
• num_predictions = B * (L-1).
• Return the three values in that order (inputs and targets as np.ndarray).

Notes

• This one-token shift is the whole trick behind next-token training — no labels needed, every position teaches the model what comes next.