How to Train a Decoder for Pre-trained BERT Transformer-Encoder?
How to Train a Decoder for a Pre‑trained BERT Transformer‑Encoder
When BERT first burst onto the NLP scene, it was celebrated for its powerful encoder‑only architecture. Since then, researchers have explored ways to pair a decoder with BERT’s encoder to create seq2seq systems for tasks such as summarization, translation, and question generation. This guide walks through the key steps for training a decoder that can effectively leverage a pre‑trained BERT encoder.
1. Choose the Right Decoder Architecture
- Transformer Decoder – Mirrors the original Transformer decoder (masked self‑attention + cross‑attention). Works well with BERT because the hidden‑size dimensions already match.
- GPT‑style Decoder – Unidirectional, autoregressive stacks of self‑attention layers. Can be initialized from a pre‑trained GPT model and then fine‑tuned with BERT’s encoder outputs.
- Lightweight RNN/Conv Decoder – Useful for low‑resource settings; however, cross‑attention must be added manually to connect to BERT.
2. Align Tokenization and Embeddings
BERT uses WordPiece tokenization and a 768‑dimensional embedding (for bert-base). Ensure your decoder shares the same tokenizer and embedding matrix, or map BERT’s embeddings to the decoder’s vocabulary with a linear projection.
# Example (PyTorch)
bert_tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
decoder_vocab = bert_tokenizer.vocab # share vocab
decoder_embeddings = nn.Embedding(len(decoder_vocab), bert_model.config.hidden_size)
decoder_embeddings.weight = bert_model.embeddings.word_embeddings.weight
3. Freeze or Fine‑Tune the Encoder?
Two common strategies:
- Freeze BERT – Keep the encoder weights fixed and only train the decoder. Faster convergence and reduces GPU memory.
- Joint Fine‑Tuning – Unfreeze selected BERT layers (often the top 4) and train encoder + decoder together. Yields better performance on downstream tasks but requires careful learning‑rate scheduling.
4. Set Up Cross‑Attention
The decoder must attend to the encoder’s hidden states. In a standard Transformer decoder block, this is the encoder‑decoder attention sub‑layer:
def decoder_block(x, enc_outputs, src_mask, tgt_mask):
# Masked self‑attention
x = self_attn(x, x, x, tgt_mask)
# Encoder‑decoder attention
x = cross_attn(x, enc_outputs, enc_outputs, src_mask)
# Feed‑forward
return feed_forward(x)
Make sure the src_mask (attention mask for the encoder) respects BERT’s padding tokens.
5. Choose an Appropriate Loss Function
- Cross‑Entropy with label smoothing (e.g., 0.1) – Standard for token‑level generation.
- Coverage Loss – Helps avoid repetitive attention when generating long sequences.
- Sequence‑level objectives (e.g., BLEU or ROUGE reinforcement learning) – Optional fine‑tuning step after initial maximum‑likelihood training.
6. Training Pipeline
- Data Preparation: Encode source sentences with BERT, obtain
enc_outputsand attention mask. - Teacher Forcing: Feed the ground‑truth target tokens shifted right into the decoder.
- Optimization: Use AdamW with separate learning rates (e.g.,
5e-5for BERT,1e-4for decoder). Apply a linear warm‑up followed by cosine decay. - Gradient Clipping: Clip at 1.0 to stabilize training.
7. Evaluation & Inference
During inference, switch BERT to evaluation mode and use beam search (size 4–8) with length penalties. Remember to mask future tokens in the decoder’s self‑attention.
8. Common Pitfalls & Tips
- Mismatched Hidden Sizes: If you use a larger BERT (e.g.,
bert-large) and a smaller decoder, add a linear projection layer between encoder outputs and decoder inputs. - Token‑Level Alignment: BERT’s WordPiece splits can cause the decoder to generate sub‑words that need to be merged post‑processing.
- Over‑fitting: When fine‑tuning the encoder, monitor validation loss closely; early stopping after 3–5 epochs often suffices.
- Memory Usage: Encoder‑decoder attention can double GPU memory. Use
torch.cuda.amp.autocast()for mixed‑precision training.
9. Sample Code Skeleton (PyTorch)
import torch
from transformers import BertModel, BertTokenizer, AdamW
# Load pre‑trained BERT encoder
bert = BertModel.from_pretrained('bert-base-uncased')
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
# Simple Transformer decoder (3 layers)
class SimpleDecoder(torch.nn.Module):
def __init__(self, hidden_dim, vocab_size, n_layers=3):
super().__init__()
self.embedding = torch.nn.Embedding(vocab_size, hidden_dim)
self.layers = torch.nn.ModuleList([
torch.nn.TransformerDecoderLayer(
d_model=hidden_dim,
nhead=8,
dim_feedforward=hidden_dim*4,
dropout=0.1)
for _ in range(n_layers)
])
self.fc_out = torch.nn.Linear(hidden_dim, vocab_size)
def forward(self, tgt_ids, enc_outputs, tgt_mask, src_mask):
tgt_emb = self.embedding(tgt_ids) * math.sqrt(self.hidden_dim)
tgt = tgt_emb.transpose(0, 1) # (seq_len, batch, dim)
memory = enc_outputs.transpose(0, 1)
for layer in self.layers:
tgt = layer(tgt, memory, tgt_mask=tgt_mask, memory_mask=src_mask)
logits = self.fc_out(tgt.transpose(0, 1))
return logits
decoder = SimpleDecoder(
hidden_dim=bert.config.hidden_size,
vocab_size=len(tokenizer)
)
# Optimizer with separate LR groups
optimizer = AdamW([
{'params': bert.parameters(), 'lr': 5e-5},
{'params': decoder.parameters(), 'lr': 1e-4}
])
# Training loop (simplified)
for batch in dataloader:
src_text, tgt_text = batch
src_ids = tokenizer(src_text, return_tensors='pt',
padding=True, truncation=True).input_ids
tgt_ids = tokenizer(tgt_text, return_tensors='pt',
padding=True, truncation=True).input_ids
# Encoder
enc_outputs = bert(input_ids=src_ids,
attention_mask=(src_ids != tokenizer.pad_token_id))[0]
# Decoder input (shifted right)
decoder_input = tgt_ids[:, :-1]
decoder_target = tgt_ids[:, 1:]
# Masks
tgt_mask = torch.nn.functional.generate_square_subsequent_mask(decoder_input.size(1))
src_mask = (src_ids != tokenizer.pad_token_id).unsqueeze(1)
# Forward
logits = decoder(decoder_input,
enc_outputs,
tgt_mask=tgt_mask.to(logits.device),
src_mask=src_mask)
loss = torch.nn.functional.cross_entropy(
logits.view(-1, logits.size(-1)),
decoder_target.view(-1),
ignore_index=tokenizer.pad_token_id,
label_smoothing=0.1
)
loss.backward()
torch.nn.utils.clip_grad_norm_(list(bert.parameters()) + list(decoder.parameters()), 1.0)
optimizer.step()
optimizer.zero_grad()
Conclusion
Training a decoder for a pre‑trained BERT encoder transforms an encoder‑only model into a full‑blown seq2seq system capable of generative tasks. The crucial ingredients are compatible tokenization, proper cross‑attention wiring, and a thoughtful training schedule that balances encoder stability with decoder flexibility. By following the steps outlined above, you can build a robust BERT‑plus‑decoder architecture tailored to your specific downstream application.
