TechByteByByte

Training a Transformer

Understand how a Transformer is trained — shifted tokens, teacher forcing, and cross-entropy loss computed at every position simultaneously — connecting directly to the loss, backpropagation, and optimizer mechanics from the Deep Learning course.

#Transformers#Training#Teacher Forcing#Cross-Entropy#AI#LLM

Begin with the central question

How does next-token error change billions of model parameters?

Essential words

A training example contains input tokens and target next tokens. Cross-entropy loss penalizes low probability on the target. Backpropagation computes gradients used by an optimizer.

What You Will Understand

How a Transformer actually gets its weights during training — the specific “shifted tokens” trick that turns raw text into a supervised learning problem, teacher forcing, and how cross-entropy loss (DL Module 6) is computed at every position simultaneously, not just the last one like Module 14’s inference trace. This module reuses your Deep Learning course’s training mechanics directly, applied specifically to Transformers.

tokens shifted into inputs/targets -> logits -> loss -> gradients -> optimizer update

The problem this module solves

Module 14 traced inference: given a prompt, predict the next token using only the last position. But training needs labeled examples — and raw text has no explicit labels. This module covers the specific, elegant trick that solves this: the text itself, shifted by one position, IS the label. No separate labeling process needed.

Build the intuition

take any sentence and create a “quiz” from it automatically: at every position, the question is “what’s the actual next word?” and the answer is simply… the actual next word, already sitting right there in the original text. No human ever has to write these answers — the training data generates its own labels, for free, at every single position in every sentence.

4. Real-World Analogy

Imagine practicing typing by re-typing a book, one word at a time, and after each word you type, immediately checking whether you typed the correct next word by looking at the book itself. The book is both your practice material and your answer key simultaneously — you never need a separate teacher grading you; the correct answer is always the very next word in the original text.

Analogy: The Auto-Quiz Practice Book & The Spelling Bee Tutor Think of shifted labels and teacher forcing as a spelling tutor correcting a child:

  • The Auto-Quiz (Shifted Labels): You print a sentence on index cards. By copying the cards and shifting them left by one spot, you instantly generate a quiz question-and-answer pair for every word: “If you see ‘the’, predict ‘cat’. If you see ‘the cat’, predict ‘sat’.”
  • The Tutor (Teacher Forcing): If a child is spelling “T-R-A-N-S-F-O-R-M-E-R” and gets the third letter wrong (saying “T-R-O”), a poor tutor would let them keep going, producing “T-R-O-N-S-F-O-R…” where every subsequent letter is now incorrect because of the early error.
  • A good tutor (Teacher Forcing) stops them instantly, replaces the wrong letter with the correct “A”, and asks: “Now that we have ‘T-R-A’, what is the next letter?”
  • This keeps training aligned and allows us to compute spelling grades for all positions simultaneously in one step.

📊 Visual Chart: Shifted Labels and Teacher Forcing Alignment

Here is how input IDs match target label IDs during parallel cross-entropy calculation:

graph TD
    classDef input fill:#2ecc71,stroke:#333,stroke-width:1px,color:#fff;
    classDef label fill:#3498db,stroke:#333,stroke-width:1px,color:#fff;
    classDef loss fill:#f1c40f,stroke:#333,stroke-width:1px,color:#fff;




    subgraph Inputs ["Input Token IDs (sees current context)"]
        In0["Position 0: '<bos>'"]:::input
        In1["Position 1: 'the'"]:::input
        In2["Position 2: 'cat'"]:::input
        In3["Position 3: 'sat'"]:::input
    end




    subgraph Labels ["Target Label IDs (correct next words)"]
        Lab0["'the'"]:::label
        Lab1["'cat'"]:::label
        Lab2["'sat'"]:::label
        Lab3["'on'"]:::label
    end




    In0 -.-> Cross0["Cross Entropy Loss 0"]:::loss
    Lab0 -.-> Cross0




    In1 -.-> Cross1["Cross Entropy Loss 1"]:::loss
    Lab1 -.-> Cross1




    In2 -.-> Cross2["Cross Entropy Loss 2"]:::loss
    Lab2 -.-> Cross2




    In3 -.-> Cross3["Cross Entropy Loss 3"]:::loss
    Lab3 -.-> Cross3




    Cross0 --> Sum["Average Loss across positions"]
    Cross1 --> Sum
    Cross2 --> Sum
    Cross3 --> Sum




    Sum --> Backprop["Backward Pass Gradients"]

5. Core Concept

Shifted tokens (the core trick):




  Original sequence:  <bos> the cat sat on the mat <eos>




  Input  (positions 0-6):  <bos> the cat sat on the mat
  Labels (positions 1-7):        the cat sat on the mat <eos>




  The labels are just the INPUT, SHIFTED by one position.
TermDefinition
Shifted tokensUsing the same sequence as both input and labels, offset by one position
Teacher forcingDuring training, the model is always shown the TRUE previous tokens (not its own, possibly-wrong, predictions) as input for predicting the next one
Training examplesEvery single POSITION in every sequence is one training example — not just one example per sentence

🧠 This is exactly your ML course’s supervised learning (features + labels), applied to text with a specific, elegant trick for generating labels automatically — precisely your Deep Learning course’s self-supervised learning concept (DL Module 2), now made completely concrete for Transformers.

6. How It Works — Step by Step

1. Take a sequence of tokens (a training document/sentence)
2. Create INPUT = the sequence, all but the last token
3. Create LABELS = the same sequence, all but the first token
   (i.e., labels[i] = input[i+1] -- the ACTUAL next token)
4. Run the model's forward pass on INPUT (Module 14's exact
   mechanism -- causal decoder blocks, LM head, softmax)
   -- but now compute logits/probabilities for EVERY position,
   not just the last one
5. At EVERY position, compare the predicted probability
   distribution against that position's TRUE label using
   cross-entropy loss (DL Module 6)
6. Average (or sum) the loss across ALL positions -- this ONE
   number is what backpropagation (DL Module 7) computes
   gradients from
7. The optimizer (DL Module 9) updates EVERY parameter using
   those gradients
8. Repeat across enormous numbers of sequences, for enormous
   numbers of steps (DL Module 8)

Teacher forcing, specifically: at position 5, the model’s input includes the true token at position 4 (from the training data), not whatever the model itself might have predicted at position 4 if it had been generating freely. This makes training dramatically more stable and parallelizable — every position’s loss can be computed simultaneously, in one forward pass, rather than needing to generate sequentially.

7. Mathematical Intuition

Read the mathematics as a story

Training shifts one token sequence into aligned inputs and targets. Every position predicts the token immediately to its right, and cross-entropy measures each mistake.

tokens:  [the, cat, sat]
inputs:  [the, cat]
targets: [cat, sat]
loss = average target surprise

The total training loss for one sequence is the sum (or average) of cross-entropy loss at every position:

total_loss = Σ CrossEntropy(predicted_distribution[i], true_label[i])
             for every position i in the sequence

This is directly your Deep Learning course’s cross-entropy (DL Module 6) — just computed once per position and combined, rather than computed once total. Every position genuinely contributes to the gradient computation simultaneously — this is precisely what teacher forcing enables: no need to wait for one position’s prediction before computing loss for the next.

8. Small Worked Example

Walk through the example

  1. Shift the sequence. 2. Produce logits for all input positions in parallel. 3. read target probabilities. 4. Calculate loss. 5. Backpropagate and update.

For the sentence “the cat sat,” at position 0 (seeing “the”), the label is “cat” — the model’s predicted probability for “cat” at this position directly determines this position’s loss.

At position 1 (seeing “the cat”), the label is “sat.” Both positions’ losses are computed from the same single forward pass — this is the direct, practical benefit of teacher forcing plus shifted labels: full parallelization across positions during training, exactly mirroring Module 1’s parallelization advantage for attention itself.

9. Python / NumPy Example

What the code will demonstrate

This small NumPy example makes Training a Transformer visible with inspectable numbers and shapes. Read it in three passes: identify each input, follow the transformation line by line, and connect the printed output to the diagram above. The arrays are intentionally tiny teaching values; unless the text explicitly says otherwise, they are not weights or measurements from GPT, Gemini, or another trained model.

# The arrays are intentionally small so each transformation can be inspected.
# Printed values illustrate the mechanism; they are not trained-model measurements.
import numpy as np




def softmax(x, axis=-1):
    exp_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return exp_x / np.sum(exp_x, axis=axis, keepdims=True)




np.random.seed(15)




vocab = ["<bos>", "the", "cat", "sat", "on", "mat", "<eos>"]
sentence = ["<bos>", "the", "cat", "sat", "on", "the", "mat", "<eos>"]
token_ids = [vocab.index(t) if t in vocab else vocab.index("the") for t in sentence]




# Input: everything EXCEPT the last token
# Labels: everything EXCEPT the first token (SHIFTED by one position)
input_ids = token_ids[:-1]
label_ids = token_ids[1:]




print("Full sequence:", sentence)
print("Input IDs (what the model SEES):   ", input_ids)
print("Label IDs (what it should PREDICT):", label_ids)
print("\nAt each position, input token -> label (correct next token):")
for i in range(len(input_ids)):
    print(f"  position {i}: sees '{sentence[i]}' -> should predict '{sentence[i+1]}'")




# --- Cross-entropy loss across ALL positions simultaneously ---
seq_len = len(input_ids)
vocab_size = len(vocab) + 1




# Simulated model logits per position (standing in for a real forward pass,
# already verified end to end in Module 14)
logits_per_position = np.random.randn(seq_len, vocab_size) * 1.5




total_loss = 0
for pos in range(seq_len):
    probs = softmax(logits_per_position[pos])
    true_id = label_ids[pos] if label_ids[pos] < vocab_size else 0
    position_loss = -np.log(probs[true_id] + 1e-10)
    total_loss += position_loss
    print(f"Position {pos}: predicted prob of true next token = {probs[true_id]:.4f}, loss = {position_loss:.4f}")




avg_loss = total_loss / seq_len
print(f"\nTotal loss across all {seq_len} positions: {total_loss:.4f}")
print(f"Average loss per position: {avg_loss:.4f}")

Expected Output:

Full sequence: ['<bos>', 'the', 'cat', 'sat', 'on', 'the', 'mat', '<eos>']
Input IDs (what the model SEES):    [0, 1, 2, 3, 4, 1, 5]
Label IDs (what it should PREDICT): [1, 2, 3, 4, 1, 5, 6]




At each position, input token -> label (correct next token):
  position 0: sees '<bos>' -> should predict 'the'
  position 1: sees 'the' -> should predict 'cat'
  position 2: sees 'cat' -> should predict 'sat'
  position 3: sees 'sat' -> should predict 'on'
  position 4: sees 'on' -> should predict 'the'
  position 5: sees 'the' -> should predict 'mat'
  position 6: sees 'mat' -> should predict '<eos>'
Position 0: predicted prob of true next token = 0.3060, loss = 1.1841
Position 1: predicted prob of true next token = 0.0666, loss = 2.7095
Position 2: predicted prob of true next token = 0.0565, loss = 2.8736
Position 3: predicted prob of true next token = 0.3504, loss = 1.0485
Position 4: predicted prob of true next token = 0.0031, loss = 5.7786
Position 5: predicted prob of true next token = 0.4342, loss = 0.8342
Position 6: predicted prob of true next token = 0.0126, loss = 4.3776




Total loss across all 7 positions: 18.8060
Average loss per position: 2.6866

10. How It Works

  • label_ids is exactly input_ids shifted by one position — verified directly: label_ids[i] == token_ids[i+1] for every position.
  • Every single position gets its own loss, computed from the same single forward pass — position 4’s very low predicted probability (0.0031, giving a high loss of 5.78) shows this (untrained, random) model badly mispredicting “the” after “on,” while position 5 did comparatively well (0.4342 probability, low loss 0.83).
  • The total_loss (or its average) is one single number — exactly what DL Module 7’s backpropagation needs as its starting point — even though it was computed by combining losses from 7 different positions simultaneously.

11. How Is This Used in Modern AI?

Where this concept lives

Follow the concept at three levels: inside the model, where the computation happens; inside the AI product, where that computation supports a visible feature; and inside production, where engineers measure speed, memory, quality, and failure cases. The details below connect those levels.

🤖 How Is This Used in Modern AI?

This shifted-tokens mechanism is precisely how every LLM’s pretraining works, applied across an enormous corpus: every sentence, in every document, at every position, becomes a free training example, with no human labeling required at all. This is exactly DL Module 2’s self-supervised learning, now completely concrete.

ConceptAI application
Shifted tokensThe core mechanism behind LLM pretraining at internet scale
Teacher forcingWhat allows a full sequence’s loss to be computed in one parallel forward pass during training
Cross-entropy per positionSummed/averaged across the entire batch and sequence to produce one training loss (DL Module 6-8)

12. How Is This Used in LLMs?

Trace one model call

User text → tokens → Transformer computation → output-token probabilities
          this topic affects one part of that computation

An LLM does not apply this idea as a separate magic step. It uses it as part of the repeated numerical pipeline that transforms token vectors and produces the next-token probabilities.

Fine-tuning (ML course Module 19, DL course Module 16) uses this exact same shifted-tokens/cross-entropy mechanism — just starting from pretrained weights instead of random initialization, on a smaller, task-specific dataset, typically with a lower learning rate (DL Module 9) to avoid overwriting the model’s existing knowledge too aggressively.

Real systems you can recognize

Transformer language models are commonly pretrained with token-prediction objectives, then may receive instruction tuning and preference optimization. Public GPT/Gemini APIs do not reveal all foundation-model training data or exact parameter counts, so the tutorial avoids invented numbers.

13. How Is This Used in Agentic AI?

Trace one agent step

Goal + history + tool results
              ↓
     LLM processes the context
              ↓
Suggested answer or tool call
              ↓
Agent runtime validates and executes it

This distinction matters: the Transformer helps produce the proposal, while the surrounding agent software controls tools, permissions, retries, memory, and execution.

Direct relevance to Agentic AI: Low-to-Moderate, indirectly. Most AI engineers building agents won’t pretrain a model from scratch — but understanding this mechanism explains why fine-tuning an LLM for a specific agent capability (e.g., a specialized tool-use style) is a real, well-understood training process, not a mysterious black box, and why it has genuine hyperparameters (learning rate, epochs, DL Module 8-9) worth configuring thoughtfully.

When this knowledge is useful

Use Training a Transformer when you need to explain, implement, debug, evaluate, or optimize the corresponding part of a Transformer pipeline. It is also useful when a model API behaves unexpectedly and you need to trace the behavior back to tokens, tensor shapes, attention visibility, training, or inference mechanics.

When it is not enough

Understanding this mechanism does not by itself prove that a complete model or application is accurate, safe, fast, or cost-effective. Production decisions still require representative evaluation data, latency and memory measurements, model-specific documentation, and tests of the surrounding retrieval or agent code.

14. Common Beginner Mistakes

⚠️ Mistake

Incorrect idea: assuming training requires separately generating output token by token, the way inference does (Module 14).

Why it is incorrect: It doesn’t — teacher forcing means the entire sequence’s loss is computed in one parallel forward pass, using the true tokens throughout, not the model’s own (possibly wrong) predictions at each step.

⚠️ Mistake

Incorrect idea: thinking only the LAST position contributes to training loss.

Why it is incorrect: Unlike inference (Module 14), where only the last position’s prediction is used, training computes and uses loss from every position simultaneously — verified directly, with all 7 positions contributing to total_loss.

⚠️ Mistake

Incorrect idea: believing labels need to be separately created by humans.

Why it is incorrect: As demonstrated, labels are simply the input sequence shifted by one position — generated automatically from the raw text itself, no human labeling involved (DL Module 2’s self-supervised learning).

15. Important Distinctions

Training (this module)Inference (Module 14)
Uses EVERY position’s prediction to compute lossUses only the LAST position’s prediction
Teacher forcing: true previous tokens as inputThe model’s own previously-generated tokens as input
Parameters UPDATE (DL Module 7-9)Parameters stay FIXED
Shifted TokensManual Labeling (classical supervised learning)
Labels generated automatically from the text itselfLabels require human annotation effort
Free, scales to any amount of raw textExpensive, limited by annotation capacity

16. Production / Engineering Considerations

  • Batch training processes many sequences simultaneously (DL Module 8’s mini-batch gradient descent), each contributing per-position losses that get combined into one batch loss.
  • Training compute cost scales with both the amount of text processed and the sequence length — since every position in every sequence contributes to the loss computation, longer training sequences mean proportionally more loss terms per forward pass.

17. Interview Questions

Beginner

Q: How does raw text become labeled training data for an LLM?

Ans: Through “shifted tokens” — the same sequence of text is used as both the input and, shifted by one position, the labels. At every position, the “label” is simply the actual next token that appears in the original text — no separate human labeling process is needed.

Intermediate

Q: What is teacher forcing, and why does it matter for training efficiency?

Ans: Teacher forcing means that during training, the model is always given the TRUE previous tokens as input for predicting the next one — not its own potentially-incorrect predictions.

This allows an entire sequence’s loss to be computed in a single, parallel forward pass (using shifted labels), rather than requiring the model to generate tokens one at a time sequentially during training, which would be far slower and less stable.

Advanced

Q: How does training’s loss computation differ from inference’s, given that both use the same underlying forward pass mechanism from Module 14?

Ans: During inference, only the last position’s predicted probability distribution is used, to select the next token to generate.

During training, EVERY position’s predicted probability distribution is compared against its corresponding true label (from the shifted-token scheme) via cross-entropy loss, and these per-position losses are combined (summed or averaged) into one total loss value — verified directly in this module, where all 7 positions in a training sequence contributed to the total training loss simultaneously, from one single forward pass.

Scenario

Q: A team is fine-tuning an LLM and notices the loss at certain positions in their training sequences (e.g., near the end of long documents) is consistently much higher than at other positions. What might this indicate?

Ans: Since every position contributes its own cross-entropy loss (as demonstrated), consistently higher loss at specific positions suggests the model is systematically struggling to predict tokens in that context — possibly indicating those positions involve less predictable content (e.g., rare terms, or content the pretrained model has less prior exposure to), or potentially a data quality issue in that portion of the training documents.

Examining which specific tokens/positions have high loss is a genuine, practical debugging technique for understanding what a model is or isn’t learning well.

Architecture

Q: Why is it accurate to say LLM pretraining is a form of self-supervised learning, connecting to your Deep Learning course?

Ans: Self-supervised learning (DL Module 2) uses labels generated automatically from the data itself, rather than requiring human annotation.

Shifted-token training is a direct, concrete instance of this: the “label” at every position — the true next token — already exists within the raw, naturally-occurring training text, requiring zero human labeling effort, yet still allowing training with the exact same supervised-style cross-entropy loss and backpropagation mechanics as any labeled classification task.

Engineering

Q: Why does training compute cost scale with sequence length, given that a single forward pass processes an entire sequence at once?

Ans: Longer sequences mean more token positions, each requiring its own attention computation (scaling with sequence length, Module 17 covers this precisely) and each contributing its own cross-entropy loss term that must be computed and backpropagated through.

Even though teacher forcing allows this to happen in one parallel forward/backward pass rather than sequentially, the total amount of computation within that pass still grows with sequence length — a genuine, practical cost driver for training on longer documents or contexts.

18. What You Should Remember

  • Shifted tokens turn raw text into free, automatically-labeled training data: labels are simply the input sequence offset by one position.
  • Teacher forcing lets an entire sequence’s loss be computed in one parallel forward pass, using true tokens as input throughout — verified directly across 7 simultaneous position losses.
  • Training uses every position’s loss; inference (Module 14) uses only the last position’s prediction — a fundamental, verified distinction.

19. How This Helps Me Build AI Systems

You now understand precisely how the weights driving Module 14’s next-token prediction trace actually get learned in the first place — using nothing but your Deep Learning course’s existing loss, backpropagation, and optimizer mechanics (DL Module 6-9), applied to text via the elegant shifted-token trick. This also directly explains what’s happening, mechanically, whenever you or your team fine-tunes a model.

Next: Module 16 — Transformer Inference and KV Cache — why inference works differently from training, and the caching technique that makes real-world LLM serving practical.

Author
TechByteByByte Editorial Team
Reviewed by
TechByteByByte Admin
Published
Last reviewed