Teacher Forcing › Code Implementation
Language AI HandbookAll chapters Teacher Forcing › Code ImplementationOne section · Sections in this chapter12 / 16 ↓
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
Let's implement teacher forcing and compare it to scheduled sampling using a simple character-level seq2seq model. This will make the mechanics concrete and reveal the exposure bias effect directly.
First, we install the necessary libraries and set up the basic data.
We will create a simple synthetic task: reversing a short sequence of integers. The encoder reads the input forward, and the decoder must output it in reverse order. This task is simple enough to train quickly but complex enough to show the teacher forcing effect. It requires the decoder to learn dependencies across the full sequence length, which means exposure bias will manifest even on short sequences.
# Simple synthetic dataset: reverse a sequence of integers
# Input: [3, 7, 5, 4, 8, 6] -> Output: [6, 8, 4, 5, 7, 3]
VOCAB_SIZE = 12
SEQ_LEN = 6
PAD_TOKEN = 0
SOS_TOKEN = 1
EOS_TOKEN = 2
def generate_data(n_samples, seq_len=SEQ_LEN, vocab_size=VOCAB_SIZE):
"""Generate (input, target) pairs where target is reversed input."""
data = []
for _ in range(n_samples):
seq = [random.randint(3, vocab_size - 1) for _ in range(seq_len)]
src = seq
tgt = [SOS_TOKEN] + list(reversed(seq)) + [EOS_TOKEN]
data.append((src, tgt))
return data
random.seed(42)
torch.manual_seed(42)
np.random.seed(42)
train_data = generate_data(2000)
val_data = generate_data(500)Training samples: 2000 Sample input: [4, 3, 7, 6, 6, 5] Sample target: [1, 5, 6, 6, 7, 3, 4, 2]
Now we define the encoder and decoder. The encoder is a simple LSTM that processes the input sequence and returns the final hidden state, which is the initial state for the decoder. The decoder is an LSTM that generates the output one token at a time, using its previous token embedding and its recurrent hidden state.
class Encoder(nn.Module):
def __init__(self, vocab_size, embed_dim, hidden_dim):
super().__init__()
self.embedding = nn.Embedding(
vocab_size, embed_dim, padding_idx=PAD_TOKEN
)
self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True)
def forward(self, x):
embedded = self.embedding(x)
_, (hidden, cell) = self.lstm(embedded)
return hidden, cell
class Decoder(nn.Module):
def __init__(self, vocab_size, embed_dim, hidden_dim):
super().__init__()
self.embedding = nn.Embedding(
vocab_size, embed_dim, padding_idx=PAD_TOKEN
)
self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True)
self.fc = nn.Linear(hidden_dim, vocab_size)
def forward(self, x, hidden, cell):
# x shape: (batch, 1) - single token
embedded = self.embedding(x)
output, (hidden, cell) = self.lstm(embedded, (hidden, cell))
prediction = self.fc(output.squeeze(1))
return prediction, hidden, cellNow we implement the seq2seq model with a configurable teacher_forcing_ratio. When set to 1.0, the decoder always receives the ground-truth previous token. When set to 0.0, it always uses its own predictions, fully replicating inference-time conditions.
class Seq2Seq(nn.Module):
def __init__(self, encoder, decoder, device):
super().__init__()
self.encoder = encoder
self.decoder = decoder
self.device = device
def forward(self, src, tgt, teacher_forcing_ratio=1.0):
batch_size = src.shape[0]
tgt_len = tgt.shape[1]
vocab_size = self.decoder.fc.out_features
outputs = torch.zeros(batch_size, tgt_len, vocab_size).to(self.device)
hidden, cell = self.encoder(src)
# First decoder input is always the SOS token
decoder_input = tgt[:, 0].unsqueeze(1)
for t in range(1, tgt_len):
pred, hidden, cell = self.decoder(decoder_input, hidden, cell)
outputs[:, t, :] = pred
# Teacher forcing: use ground truth or model prediction?
if random.random() < teacher_forcing_ratio:
decoder_input = tgt[:, t].unsqueeze(1) # ground truth
else:
decoder_input = pred.argmax(1).unsqueeze(1) # model prediction
return outputsEMBED_DIM = 16
HIDDEN_DIM = 64
BATCH_SIZE = 64
LEARNING_RATE = 0.001
N_EPOCHS = 30
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def make_model():
encoder = Encoder(VOCAB_SIZE, EMBED_DIM, HIDDEN_DIM)
decoder = Decoder(VOCAB_SIZE, EMBED_DIM, HIDDEN_DIM)
model = Seq2Seq(encoder, decoder, device).to(device)
return model
def make_batch(data, batch_size=BATCH_SIZE):
"""Sample a random minibatch and convert to tensors."""
batch = random.sample(data, batch_size)
src = torch.tensor([b[0] for b in batch], dtype=torch.long).to(device)
tgt = torch.tensor([b[1] for b in batch], dtype=torch.long).to(device)
return src, tgt
def train_epoch(model, data, optimizer, criterion, teacher_forcing_ratio):
model.train()
total_loss = 0.0
n_batches = len(data) // BATCH_SIZE
for _ in range(n_batches):
src, tgt = make_batch(data)
optimizer.zero_grad()
output = model(src, tgt, teacher_forcing_ratio)
# Flatten predictions and targets, skipping the SOS position
output_flat = output[:, 1:, :].reshape(-1, VOCAB_SIZE)
tgt_flat = tgt[:, 1:].reshape(-1)
loss = criterion(output_flat, tgt_flat)
loss.backward()
optimizer.step()
total_loss += loss.item()
return total_loss / n_batches
def evaluate(model, data, criterion, teacher_forcing_ratio=0.0):
"""Evaluate using free-running decoding to measure inference quality."""
model.eval()
total_loss = 0.0
n_batches = len(data) // BATCH_SIZE
with torch.no_grad():
for _ in range(n_batches):
src, tgt = make_batch(data)
output = model(src, tgt, teacher_forcing_ratio)
output_flat = output[:, 1:, :].reshape(-1, VOCAB_SIZE)
tgt_flat = tgt[:, 1:].reshape(-1)
loss = criterion(output_flat, tgt_flat)
total_loss += loss.item()
return total_loss / n_batchesNow we train two models side by side: one with pure teacher forcing, and one with scheduled sampling that linearly decays from full teacher forcing to a 50/50 mix. Both models have identical architecture and initialization (same model factory, same random seeds). The only difference is the training strategy.
criterion = nn.CrossEntropyLoss(ignore_index=PAD_TOKEN)
# Model 1: pure teacher forcing
model_tf = make_model()
initial_model_state = copy.deepcopy(model_tf.state_dict())
optimizer_tf = optim.Adam(model_tf.parameters(), lr=LEARNING_RATE)
tf_train_losses = []
tf_val_losses = []
for epoch in range(N_EPOCHS):
train_loss = train_epoch(
model_tf, train_data, optimizer_tf, criterion, teacher_forcing_ratio=1.0
)
val_loss = evaluate(
model_tf, val_data, criterion, teacher_forcing_ratio=0.0
)
tf_train_losses.append(train_loss)
tf_val_losses.append(val_loss)
# Model 2: scheduled sampling (linear decay from 1.0 to 0.5)
model_ss = make_model()
model_ss.load_state_dict(initial_model_state)
optimizer_ss = optim.Adam(model_ss.parameters(), lr=LEARNING_RATE)
ss_train_losses = []
ss_val_losses = []
for epoch in range(N_EPOCHS):
# Linear schedule: 1.0 at epoch 0, 0.5 at final epoch
tf_ratio = max(0.5, 1.0 - epoch * (0.5 / max(1, N_EPOCHS - 1)))
train_loss = train_epoch(
model_ss,
train_data,
optimizer_ss,
criterion,
teacher_forcing_ratio=tf_ratio,
)
val_loss = evaluate(
model_ss, val_data, criterion, teacher_forcing_ratio=0.0
)
ss_train_losses.append(train_loss)
ss_val_losses.append(val_loss)Pure teacher forcing - final free-run val loss: 0.1281 Scheduled sampling - final free-run val loss: 0.1173 Scheduled sampling is 0.0107 lower loss (+0.0107)
The validation loss here is always measured with free-running decoding (no teacher forcing), which mirrors actual inference conditions. Notice the difference between training loss (which is computed on teacher-forced inputs) and this validation metric. A model can have low training loss while maintaining high free-running validation loss: this gap is the measurable signature of exposure bias. The scheduled sampling model's training loss may be somewhat higher because it is learning the harder problem of operating on imperfect inputs, but its free-running validation loss should be lower because it has practiced this mode during training.
Now let's visualize the training dynamics to see how the two strategies differ over the course of training.
Let's also implement and visualize the three standard scheduled sampling probability schedules. Seeing the shapes of these curves concretely helps build intuition for the pacing they impose on the curriculum.
def linear_schedule(i, k=1.0, c=0.02, epsilon_min=0.1):
"""Linearly decay teacher forcing probability."""
return max(epsilon_min, k - i * c)
def exponential_schedule(i, k=0.98):
"""Exponentially decay teacher forcing probability."""
return k**i
def inverse_sigmoid_schedule(i, k=10.0):
"""S-shaped (inverse sigmoid) decay schedule."""
return k / (k + np.exp(i / k))
n_steps = 100
steps = np.arange(n_steps)
linear_probs = [linear_schedule(i) for i in steps]
exp_probs = [exponential_schedule(i) for i in steps]
isigmoid_probs = [inverse_sigmoid_schedule(i) for i in steps]Finally, let's compute exact sequence match accuracy under free-running decoding for both models. Loss values can be hard to interpret; sequence accuracy tells us directly how often the model gets the full reversal correct at inference time.
def sequence_accuracy(model, data, n_samples=500):
"""Compute exact sequence match accuracy using free-running decoding."""
model.eval()
correct = 0
total = min(n_samples, len(data))
samples = data[:total]
with torch.no_grad():
for src_seq, tgt_seq in samples:
src = torch.tensor([src_seq], dtype=torch.long).to(device)
tgt = torch.tensor([tgt_seq], dtype=torch.long).to(device)
output = model(src, tgt, teacher_forcing_ratio=0.0)
predicted = output[0, 1:, :].argmax(dim=-1).cpu().tolist()
target = tgt_seq[1:] # skip SOS
if predicted == target:
correct += 1
return correct / total
acc_tf = sequence_accuracy(model_tf, val_data)
acc_ss = sequence_accuracy(model_ss, val_data)Sequence accuracy (free-run decoding): Pure teacher forcing: 94.8% Scheduled sampling: 95.6% Difference: +0.8 percentage points (scheduled sampling)
The exact sequence match accuracy under free-running decoding reveals the practical impact of exposure bias. The task requires the model to get every single token right to score a match, making it a stringent measure of inference quality. The scheduled sampling model, having practiced handling its own outputs during training, tends to produce more consistent complete sequences at inference time, especially for sequences where an early error would cascade into multiple downstream mistakes.
Key Parameters
Understanding the key hyperparameters for teacher forcing and scheduled sampling helps you apply these techniques effectively:
- teacher_forcing_ratio: Probability of using the ground-truth token as decoder input. Set to 1.0 for pure teacher forcing, 0.0 for free-running. Scheduled sampling decays this value during training. The initial value should almost always start at 1.0 to ensure stable early training.
- schedule type: The decay function controlling how teacher forcing probability decreases. Linear decay is simple and predictable; inverse sigmoid provides a more controlled transition that preserves the stable learning phase.
- epsilon_min: The floor value for scheduled sampling probability. Keeping this above 0 ensures some ground-truth guidance persists even late in training, which can stabilize learning. A common choice is 0.0 to 0.2 depending on task difficulty.
- decay rate: Controls how quickly teacher forcing probability falls. Faster decay exposes the model to its own errors sooner, which can reduce exposure bias but also destabilizes early training if the model has not yet learned reliable basic patterns.