How to Train a Decoder for Pre-trained BERT Transformer-Encoder?

How to Train a Decoder for a Pre-trained BERT Transformer Encoder

Since its release, BERT (Bidirectional Encoder Representations from Transformers) has become the go‑to backbone for a myriad of natural language processing (NLP) tasks. While BERT excels as an encoder, many downstream applications—such as text generation, summarization, and conversational agents—require a decoder component. This post walks through the practical steps of attaching and training a decoder on top of a frozen or fine‑tuned BERT encoder.

1. Why Pair BERT with a Decoder?

• Bidirectional Contextualization: BERT captures deep bidirectional information, providing rich token embeddings.
• Separation of Concerns: Keep the heavy pre‑training effort of the encoder intact while customizing the generation side for a specific task.
• Data Efficiency: Leveraging a pre‑trained encoder reduces the amount of task‑specific data needed for the decoder.

2. Architectural Choices

There are three common ways to integrate a decoder with BERT:

1. Seq2Seq with BERT Encoder + Transformer Decoder: The classic encoder–decoder setup (e.g., BERT‑Encoder + GPT‑style decoder).
2. Encoder–Decoder Fusion (BERT2BERT): Use two BERT models—one as encoder, one as decoder—sharing token embeddings.
3. Conditional Language Model: Append a special <sep> token and let the decoder attend to encoder hidden states directly (as in T5).

3. Preparing the Data

For a decoder‑centric task, you need source–target pairs:

• Source: Input sequence fed to BERT (e.g., a passage, a question, or a prompt).
• Target: Desired output sequence for the decoder (e.g., summary, answer, or generated text).

Typical preprocessing steps:

1. Tokenize both source and target with the same WordPiece/BPE vocabulary used by BERT.
2. Add special tokens: [CLS] at the start of the source, [SEP] to separate source and target, and [PAD] for padding.
3. Create attention masks for encoder and decoder, and a cross‑attention mask that prevents the decoder from attending to future tokens.

4. Implementing the Model

Below is a high‑level PyTorch‑style pseudo‑code that illustrates the wiring of a BERT encoder with a Transformer decoder.

import torch
from transformers import BertModel, BertTokenizer, BertConfig
from torch.nn import TransformerDecoder, TransformerDecoderLayer

# Load pre‑trained BERT encoder
bert_encoder = BertModel.from_pretrained('bert-base-uncased')
bert_encoder.eval()  # freeze if you don’t want to fine‑tune

# Decoder configuration – match hidden size & number of heads
config = BertConfig(
    hidden_size=bert_encoder.config.hidden_size,
    num_hidden_layers=6,          # you can choose depth
    num_attention_heads=bert_encoder.config.num_attention_heads,
    intermediate_size=bert_encoder.config.intermediate_size,
    vocab_size=bert_encoder.config.vocab_size,
)

# Create the decoder stack
decoder_layer = TransformerDecoderLayer(
    d_model=config.hidden_size,
    nhead=config.num_attention_heads,
    dim_feedforward=config.intermediate_size,
    dropout=0.1,
    activation='gelu'
)
decoder = TransformerDecoder(decoder_layer, num_layers=config.num_hidden_layers)

# Output projection to vocab
lm_head = torch.nn.Linear(config.hidden_size, config.vocab_size, bias=False)

def forward(src_ids, src_mask, tgt_ids, tgt_mask):
    # Encoder forward pass (no gradient if frozen)
    encoder_outputs = bert_encoder(input_ids=src_ids, attention_mask=src_mask)
    memory = encoder_outputs.last_hidden_state   # shape: (B, S_src, H)

    # Prepare tgt embeddings (share BERT token embeddings)
    tgt_embeddings = bert_encoder.embeddings(tgt_ids)  # (B, S_tgt, H)

    # Decoder forward pass
    decoder_output = decoder(
        tgt=tgt_embeddings.transpose(0,1),   # (S_tgt, B, H)
        memory=memory.transpose(0,1),       # (S_src, B, H)
        tgt_mask=tgt_mask,                  # subsequent mask
        memory_key_padding_mask=~src_mask.bool()
    )
    decoder_output = decoder_output.transpose(0,1)  # (B, S_tgt, H)

    # Language modeling head
    logits = lm_head(decoder_output)  # (B, S_tgt, vocab_size)
    return logits

5. Training Procedure

1. Loss Function: Use CrossEntropyLoss with ignore_index=tokenizer.pad_token_id. Apply it token‑wise on the decoder logits.
2. Optimization:
   • If the encoder is frozen, only update decoder parameters.
   • If fine‑tuning, set a lower learning rate for the encoder (e.g., lr_encoder = 1e-5) and a higher one for the decoder (e.g., lr_decoder = 5e-4).
3. Training Loop Sketch:

optimizer = torch.optim.AdamW([
    {'params': decoder.parameters()},
    {'params': bert_encoder.parameters(), 'lr': 1e-5}
], lr=5e-4)

criterion = torch.nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)

for epoch in range(num_epochs):
    for batch in dataloader:
        src_ids, src_mask, tgt_ids, tgt_mask = batch

        # Shift target for teacher forcing
        decoder_input = tgt_ids[:, :-1]
        decoder_target = tgt_ids[:, 1:]

        logits = model(src_ids, src_mask, decoder_input, tgt_mask[:, :-1])
        loss = criterion(logits.reshape(-1, vocab_size), decoder_target.reshape(-1))
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

6. Tips & Tricks

• Warm‑up & Scheduler: Use a linear warm‑up (10% of total steps) followed by cosine decay for stable training.
• Label Smoothing: Improves generalization for generation tasks.
• Gradient Checkpointing: Saves memory when fine‑tuning the encoder.
• Mixed Precision (AMP): Speeds up training on modern GPUs.
• Evaluation Metrics: BLEU, ROUGE, or METEOR depending on the task; also track perplexity for language modeling quality.

7. Common Pitfalls

1. Token Mismatch: Ensure the decoder uses the exact same tokenizer/vocabulary as the BERT encoder.
2. Masking Errors: Forgetting the causal mask for the decoder leads to “cheating” during training.
3. Over‑fitting the Decoder: If the encoder is frozen, the decoder can quickly overfit on small datasets—use dropout and early stopping.
4. Sequence Lengths: Align max source and target lengths to the model’s positional embeddings; otherwise, you’ll hit index errors.

8. Example Use Cases

• Abstractive Summarization: Encode a news article with BERT, decode a concise summary.
• Question Answering (Generation): Encode the context + question, generate a free‑form answer.
• Dialogue Systems: Encode user utterance, decode the next system response.

9. Future Directions

While the BERT‑encoder + Transformer‑decoder combo works well, newer architectures such as Encoder‑Decoder Transformers (e.g., T5, BART) train both sides jointly from scratch. Researchers are also exploring parameter-efficient tuning (adapters, LoRA) to adapt the encoder without full fine‑tuning, which can be combined with a lightweight decoder for rapid deployment.

Conclusion

Training a decoder on top of a pre‑trained BERT encoder blends the strengths of bidirectional contextual understanding with generative capabilities. By carefully wiring the encoder‑decoder attention, handling masks, and employing a balanced training schedule, you can turn BERT into a powerful backbone for any text‑generation task.