Welcome back to Bringing Transformers to Life: Training & Inference! This is an extraordinary milestone as you reach the final lesson of this course. Throughout our journey together, you've assembled a complete Transformer architecture, built robust data preparation pipelines, and implemented sophisticated training procedures with teacher forcing and learning rate scheduling. You should be incredibly proud of the deep understanding you've developed of these powerful models.
Today, we shift our focus from training to inference: the art of generating sequences with your trained Transformer model. This is where the magic truly happens, as we watch our model generate coherent text one token at a time. We'll explore two fundamental inference strategies: greedy decoding for fast generation and beam search for higher-quality output. You'll learn to implement both approaches, understand their trade-offs, and see how they perform on practical examples. By the end of this lesson, you'll have a complete inference pipeline that can generate sequences from any trained Transformer model.
From Training to Inference: A Different Challenge
Greedy Decoding: The Simplest Strategy
Beam Search: Exploring Multiple Paths
Setting Up the Inference Pipeline
The TransformerInference class encapsulates our inference functionality and provides a clean interface for generating sequences:
class TransformerInference: def __init__(self, model, src_vocab, tgt_vocab): self.model = model self.src_vocab = src_vocab self.tgt_vocab = tgt_vocab self.model.eval()
This initialization is straightforward but crucial. We store references to the trained model and both vocabularies, then call model.eval() to disable dropout and batch normalization training behaviors. This ensures consistent inference behavior and prevents the randomness that would occur during training mode, which is essential for reproducible results.
Let's examine the complete pipeline that trains a model and tests both inference strategies:
def main(): """Test inference with trained model""" print("Testing Transformer Inference...") # Prepare data and train model src_sentences, tgt_sentences = create_synthetic_data(num_samples=200) src_vocab = Vocabulary() tgt_vocab = Vocabulary() src_vocab.build_vocab(src_sentences) tgt_vocab.build_vocab(tgt_sentences) # Create and train model dataset = TranslationDataset(src_sentences, tgt_sentences, src_vocab, tgt_vocab, max_len=15) train_loader = DataLoader(dataset, batch_size=8, shuffle=True, collate_fn=collate_fn) model = Transformer( src_vocab_size=src_vocab.size, tgt_vocab_size=tgt_vocab.size, d_model=64, num_heads=4, num_encoder_layers=2, num_decoder_layers=2, d_ff=256, dropout=0.1 ) # Quick training trainer = TransformerTrainer(model, train_loader, lr=1e-3, warmup_steps=25) print("Quick training for 2 epochs...") for epoch in range(2): avg_loss = trainer.train_epoch() print(f"Epoch {epoch + 1}, Loss: {avg_loss:.4f}")
This pipeline demonstrates the complete workflow from data preparation through training to inference testing. We create synthetic data using our word reversal task, build vocabularies, train a compact Transformer model for three epochs, and prepare it for inference evaluation. The model architecture is intentionally small to enable quick training while still demonstrating meaningful learning behavior.
Testing and Analyzing Results
The inference testing reveals fascinating insights about the different strategies:
When we run this complete pipeline, we observe the following results:
Testing Transformer Inference...Quick training for 2 epochs...Step 10, Loss: 2.1828, LR: 0.011000Step 20, Loss: 1.6475, LR: 0.021000Epoch 1, Loss: 1.5726Step 35, Loss: 1.3399, LR: 0.020833Step 45, Loss: 1.3126, LR: 0.018430Epoch 2, Loss: 1.2247Step 60, Loss: 1.2136, LR: 0.016005Step 70, Loss: 1.2314, LR: 0.014835Epoch 3, Loss: 1.2724Inference Results:Source: 'hello world'Expected: 'world hello'Greedy: 'world hello' (time: 0.501s)Beam: 'world hello' (time: 3.298s)--------------------------------------------------Source: 'good morning'Expected: 'morning good'Greedy: 'morning good' (time: 0.301s)Beam: 'morning good' (time: 0.500s)--------------------------------------------------Source: 'thank you very much'Expected: 'much very you thank'Greedy: 'you to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to' (time: 3.199s)Beam: 'much very you thank' (time: 3.300s)--------------------------------------------------
The output reveals crucial insights about the different decoding strategies. For simple inputs like "hello world" and "good morning," both strategies produce perfect results, but greedy decoding is significantly faster (0.3 - 0.5 seconds vs. 0.5 - 3.3 seconds). However, the third example, "thank you very much," shows where greedy decoding fails catastrophically: it gets stuck in a repetitive loop, generating "you to to to..." until reaching the maximum length limit.
This repetitive behavior occurs because greedy decoding can become trapped in local probability maxima:
The model generates "you";
The sequence context makes "to" the most probable next token;
With "you to" in the context, the model again predicts "to" as most likely, creating a feedback loop where the model reinforces its own poor predictions.
This demonstrates why greedy decoding's myopic approach—never reconsidering previous choices—can lead to degraded output quality for complex sequences.
Conclusion and Next Steps
Congratulations on completing the final lesson of Bringing Transformers to Life: Training & Inference! You have successfully implemented a complete inference pipeline that showcases the practical application of trained Transformer models. Through greedy decoding and beam search, you've learned to balance the trade-offs between speed and quality in sequence generation, understanding when each strategy is most appropriate. Your journey from building Transformers from scratch to implementing sophisticated inference strategies represents a remarkable achievement in mastering these powerful models.
You should be incredibly proud of reaching this milestone. You've built a Transformer from the ground up, created robust data pipelines, implemented sophisticated training procedures, and now mastered the art of inference. This comprehensive understanding positions you perfectly for the next course in our learning path: Harnessing Transformers with Hugging Face, where we'll dive into the modern Hugging Face ecosystem and learn to leverage state-of-the-art pre-trained models for real-world applications. The upcoming practice exercises will solidify your understanding of inference strategies before you embark on this exciting next chapter.
Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal
Inference presents fundamentally different challenges compared to training. During training, we used teacher forcing, where the model always sees the correct previous tokens. During inference, however, the model must generate sequences autoregressively, using its own predictions as input for subsequent tokens. This creates a sequential dependency in which each prediction influences all future predictions.
The inference process follows this mathematical formulation: given a source sequence x and previously generated tokens y1,y2,...,yt−1, we compute:
P(yt∣y<t,x)=softmax(fθ(y<t,x))t
where fθ represents our trained Transformer model. We start with a special <SOS> token, generate the most likely next token, append it to the sequence, and repeat until we encounter an <EOS> token or reach a maximum length. This autoregressive nature means that early mistakes can propagate through the entire sequence, making the choice of decoding strategy crucial for output quality.
Greedy decoding represents the most straightforward approach to sequence generation: at each step, we simply select the token with the highest probability. Mathematically, this means selecting:
yt=argw∈VmaxP(w∣y<t,x)
where V is our vocabulary. While this strategy is fast and deterministic, it can lead to suboptimal sequences because it never reconsiders previous choices.
Let's implement the greedy decoding method:
def greedy_decode(self, src_sentence, max_len=50): """Greedy decoding - select most likely token at each step""" # Encode source sentence with consistent max_len src_tokens = torch.tensor([self.src_vocab.encode(src_sentence, max_len=15)], dtype=torch.long) src_mask = self.model.create_padding_mask(src_tokens) # Initialize decoder input with SOS token tgt_tokens = torch.tensor([[self.tgt_vocab.token2idx['<SOS>']]], dtype=torch.long) for _ in range(max_len): # Create both causal and padding masks for target tgt_causal_mask = self.model.create_causal_mask(tgt_tokens.size(1)) tgt_padding_mask = self.model.create_padding_mask(tgt_tokens) tgt_mask = tgt_causal_mask & tgt_padding_mask # Forward pass with no_grad for better performance with torch.no_grad(): output = self.model(src_tokens, tgt_tokens, src_mask, tgt_mask) # Get next token next_token_logits = output[:, -1, :] next_token = torch.argmax(next_token_logits, dim=-1).unsqueeze(1) # Stop if EOS token generated if next_token.item() == self.tgt_vocab.token2idx['<EOS>']: break # Append to sequence tgt_tokens = torch.cat([tgt_tokens, next_token], dim=1) return self.tgt_vocab.decode(tgt_tokens[0].tolist())
This implementation demonstrates the core autoregressive loop through the following key steps:
Source preparation: Encode the source sentence with consistent max_len and create appropriate padding masks.
Target initialization: Start with just the <SOS> token in the target sequence.
Iterative generation loop: For each step up to max_len:
Create both causal and padding masks for the current target sequence;
Perform a forward pass through the model using torch.no_grad() for better performance;
Extract the logits for the next token position and select the highest probability token with torch.argmax;
Check if the generated token is <EOS> and terminate if sequence is complete
Append the new token to the growing target sequence;
The torch.no_grad() context improves performance by disabling gradient computation during inference, which is unnecessary and computationally expensive since we're not training the model.
Beam search offers a more sophisticated approach by maintaining multiple candidate sequences (called the "beam") and exploring several possibilities simultaneously. Instead of committing to a single choice at each step, beam search keeps track of the top-k most promising sequences and expands each one. The score for each sequence is the cumulative log probability:
score(y1,...,yt)=i=1∑tlogP(yi∣y<i,x)
Here's the beam search implementation:
def beam_search(self, src_sentence, beam_width=3, max_len=50): """Beam search decoding - maintain multiple hypotheses""" # Encode source with consistent max_len src_tokens = torch.tensor([self.src_vocab.encode(src_sentence, max_len=15)], dtype=torch.long) src_mask = self.model.create_padding_mask(src_tokens) # Initialize beam with SOS token beams = [{'tokens': [self.tgt_vocab.token2idx['<SOS>']], 'score': 0.0}] for step in range(max_len): candidates = [] for beam in beams: if beam['tokens'][-1] == self.tgt_vocab.token2idx['<EOS>']: candidates.append(beam) continue # Get current sequence tgt_tokens = torch.tensor([beam['tokens']], dtype=torch.long) tgt_causal_mask = self.model.create_causal_mask(tgt_tokens.size(1)) tgt_padding_mask = self.model.create_padding_mask(tgt_tokens) tgt_mask = tgt_causal_mask & tgt_padding_mask # Forward pass with no_grad for better performance with torch.no_grad(): output = self.model(src_tokens, tgt_tokens, src_mask, tgt_mask) # Get probabilities for next token logits = output[:, -1, :] probs = F.log_softmax(logits, dim=-1) # Get top-k candidates top_probs, top_indices = torch.topk(probs, beam_width) for prob, idx in zip(top_probs[0], top_indices[0]): new_tokens = beam['tokens'] + [idx.item()] new_score = beam['score'] + prob.item() candidates.append({'tokens': new_tokens, 'score': new_score}) # Select top beam_width candidates candidates.sort(key=lambda x: x['score'], reverse=True) beams = candidates[:beam_width] # Check if all beams ended if all(beam['tokens'][-1] == self.tgt_vocab.token2idx['<EOS>'] for beam in beams): break # Return best sequence best_beam = max(beams, key=lambda x: x['score']) return self.tgt_vocab.decode(best_beam['tokens'])
Beam search maintains a list of candidate sequences with their accumulated log probability scores. This implementation demonstrates the core beam search algorithm through the following key steps:
Beam initialization: Start with a single beam containing just the <SOS> token with score 0.0.
Iterative expansion loop: For each step up to max_len:
Beam processing: Iterate through current beams, skipping those that have already terminated with <EOS>;
Forward pass: For each active beam, prepare target tokens and masks, then perform a forward pass using torch.no_grad();
Probability calculation: Extract logits and apply log_softmax to get log probabilities, avoiding numerical underflow;
Candidate generation: Use torch.topk to get the top-k most probable next tokens for each beam;
Beam expansion: Create new candidate sequences by extending each active beam with its top-k tokens;
Beam selection: Sort all candidates by accumulated score and select the top beam_width sequences;
Termination: Continue until all beams end with <EOS> tokens or maximum length is reached.
Result selection: Return the sequence with the highest accumulated log probability score.
The key insight is using log_softmax to get log probabilities, which we accumulate by addition rather than multiplication, avoiding numerical underflow issues. After expanding all beams, the algorithm selects the top beam_width candidates based on their scores and continues until all beams terminate with <EOS> tokens.
Python
class TransformerInference: def __init__(self, model, src_vocab, tgt_vocab): self.model = model self.src_vocab = src_vocab self.tgt_vocab = tgt_vocab self.model.eval()
Python
def main(): """Test inference with trained model""" print("Testing Transformer Inference...") # Prepare data and train model src_sentences, tgt_sentences = create_synthetic_data(num_samples=200) src_vocab = Vocabulary() tgt_vocab = Vocabulary() src_vocab.build_vocab(src_sentences) tgt_vocab.build_vocab(tgt_sentences) # Create and train model dataset = TranslationDataset(src_sentences, tgt_sentences, src_vocab, tgt_vocab, max_len=15) train_loader = DataLoader(dataset, batch_size=8, shuffle=True, collate_fn=collate_fn) model = Transformer( src_vocab_size=src_vocab.size, tgt_vocab_size=tgt_vocab.size, d_model=64, num_heads=4, num_encoder_layers=2, num_decoder_layers=2, d_ff=256, dropout=0.1 ) # Quick training trainer = TransformerTrainer(model, train_loader, lr=1e-3, warmup_steps=25) print("Quick training for 2 epochs...") for epoch in range(2): avg_loss = trainer.train_epoch() print(f"Epoch {epoch + 1}, Loss: {avg_loss:.4f}")
Testing Transformer Inference...Quick training for 2 epochs...Step 10, Loss: 2.1828, LR: 0.011000Step 20, Loss: 1.6475, LR: 0.021000Epoch 1, Loss: 1.5726Step 35, Loss: 1.3399, LR: 0.020833Step 45, Loss: 1.3126, LR: 0.018430Epoch 2, Loss: 1.2247Step 60, Loss: 1.2136, LR: 0.016005Step 70, Loss: 1.2314, LR: 0.014835Epoch 3, Loss: 1.2724Inference Results:Source: 'hello world'Expected: 'world hello'Greedy: 'world hello' (time: 0.501s)Beam: 'world hello' (time: 3.298s)--------------------------------------------------Source: 'good morning'Expected: 'morning good'Greedy: 'morning good' (time: 0.301s)Beam: 'morning good' (time: 0.500s)--------------------------------------------------Source: 'thank you very much'Expected: 'much very you thank'Greedy: 'you to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to to' (time: 3.199s)Beam: 'much very you thank' (time: 3.300s)--------------------------------------------------
Python
def greedy_decode(self, src_sentence, max_len=50): """Greedy decoding - select most likely token at each step""" # Encode source sentence with consistent max_len src_tokens = torch.tensor([self.src_vocab.encode(src_sentence, max_len=15)], dtype=torch.long) src_mask = self.model.create_padding_mask(src_tokens) # Initialize decoder input with SOS token tgt_tokens = torch.tensor([[self.tgt_vocab.token2idx['<SOS>']]], dtype=torch.long) for _ in range(max_len): # Create both causal and padding masks for target tgt_causal_mask = self.model.create_causal_mask(tgt_tokens.size(1)) tgt_padding_mask = self.model.create_padding_mask(tgt_tokens) tgt_mask = tgt_causal_mask & tgt_padding_mask # Forward pass with no_grad for better performance with torch.no_grad(): output = self.model(src_tokens, tgt_tokens, src_mask, tgt_mask) # Get next token next_token_logits = output[:, -1, :] next_token = torch.argmax(next_token_logits, dim=-1).unsqueeze(1) # Stop if EOS token generated if next_token.item() == self.tgt_vocab.token2idx['<EOS>']: break # Append to sequence tgt_tokens = torch.cat([tgt_tokens, next_token], dim=1) return self.tgt_vocab.decode(tgt_tokens[0].tolist())
Python
def beam_search(self, src_sentence, beam_width=3, max_len=50): """Beam search decoding - maintain multiple hypotheses""" # Encode source with consistent max_len src_tokens = torch.tensor([self.src_vocab.encode(src_sentence, max_len=15)], dtype=torch.long) src_mask = self.model.create_padding_mask(src_tokens) # Initialize beam with SOS token beams = [{'tokens': [self.tgt_vocab.token2idx['<SOS>']], 'score': 0.0}] for step in range(max_len): candidates = [] for beam in beams: if beam['tokens'][-1] == self.tgt_vocab.token2idx['<EOS>']: candidates.append(beam) continue # Get current sequence tgt_tokens = torch.tensor([beam['tokens']], dtype=torch.long) tgt_causal_mask = self.model.create_causal_mask(tgt_tokens.size(1)) tgt_padding_mask = self.model.create_padding_mask(tgt_tokens) tgt_mask = tgt_causal_mask & tgt_padding_mask # Forward pass with no_grad for better performance with torch.no_grad(): output = self.model(src_tokens, tgt_tokens, src_mask, tgt_mask) # Get probabilities for next token logits = output[:, -1, :] probs = F.log_softmax(logits, dim=-1) # Get top-k candidates top_probs, top_indices = torch.topk(probs, beam_width) for prob, idx in zip(top_probs[0], top_indices[0]): new_tokens = beam['tokens'] + [idx.item()] new_score = beam['score'] + prob.item() candidates.append({'tokens': new_tokens, 'score': new_score}) # Select top beam_width candidates candidates.sort(key=lambda x: x['score'], reverse=True) beams = candidates[:beam_width] # Check if all beams ended if all(beam['tokens'][-1] == self.tgt_vocab.token2idx['<EOS>'] for beam in beams): break # Return best sequence best_beam = max(beams, key=lambda x: x['score']) return self.tgt_vocab.decode(best_beam['tokens'])
