Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Weight Tying › Encoder-Decoder Weight Tying

Michael BrenndoerferUpdated June 17, 20254 min read

Language AI HandbookAll chapters

Part 16 of 67Transformer Architectures

Chapter 115Weight Tying › Encoder-Decoder Weight Tying

Reading section 8 of 15 · Encoder-Decoder Weight Tying · 4 minSection 8 of 15 · 4 min

Section 8 of 15 · 4 min read

Weight Tying › Encoder-Decoder Weight TyingOne section · Sections in this chapter8 / 15 ↓

Sections in this chapter15

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

So far we have discussed tying input and output embeddings within a single decoder-only model. Encoder-decoder architectures offer additional tying opportunities that can eliminate even more redundant parameters. These architectures, which underpin models like T5, BART, and the original transformer, have three embedding matrices where a decoder-only model has two.

In a sequence-to-sequence model, you have three distinct embedding operations: the encoder processes source tokens through its own embedding layer, the decoder processes target tokens through its own embedding layer during training (teacher forcing), and the decoder produces output logits through the output projection. All three matrices have the same shape $(V \times d)$ if the model uses a shared source-target vocabulary, which is common for tasks like machine translation where modern BPE tokenizers can cover both languages.

The question is whether to tie any or all of these three matrices. Research has explored several configurations:

Decoder-only tying: Tie decoder input and output, leave encoder separate
Full tying: Tie all three together with a single shared matrix
No tying: Keep all three separate

Research has shown that full tying works well in many settings:

\mathbf{E}_{\text{enc}} = \mathbf{E}_{\text{dec}} = \mathbf{W}_{\text{out}}

where:

$\mathbf{E}_{\text{enc}}$ : the encoder input embedding matrix of shape $V \times d$ , used when the encoder reads the source sequence
$\mathbf{E}_{\text{dec}}$ : the decoder input embedding matrix of shape $V \times d$ , used when the decoder reads target tokens during teacher forcing
$\mathbf{W}_{\text{out}}$ : the decoder output projection matrix of shape $V \times d$ , used when the decoder generates output logits
$V$ : the shared vocabulary size
$d$ : the hidden dimension

This three-way tying means a single learned embedding serves all three roles: encoding source tokens, encoding target tokens during teacher forcing, and defining the output distribution over the vocabulary. The parameter savings are now three-fold compared to having no tying at all.

The theoretical justification extends naturally. If source and target sequences are in the same language or share significant vocabulary (as they do in paraphrase generation, summarization, or code refinement tasks), then a token should have a consistent representation regardless of whether it appears in the input or output. Even in cross-lingual tasks, if the source and target languages share a BPE vocabulary with many overlapping subword units, some degree of tying can still be beneficial.

In[20]:

Code

class EncoderDecoderEmbeddings(nn.Module):
    """
    Shared embeddings for encoder-decoder architecture.

    All three embedding matrices are tied together.
    """

    def __init__(self, vocab_size, hidden_dim, tie_all=True):
        super().__init__()
        self.tie_all = tie_all

        # Shared embedding (used by all three roles)
        self.shared_embedding = nn.Embedding(vocab_size, hidden_dim)

        if not tie_all:
            # Separate embeddings if not tying
            self.encoder_embedding = nn.Embedding(vocab_size, hidden_dim)
            self.decoder_embedding = nn.Embedding(vocab_size, hidden_dim)
            self.output_projection = nn.Linear(
                hidden_dim, vocab_size, bias=False
            )

    def encode_input(self, token_ids):
        """Embed encoder input tokens."""
        if self.tie_all:
            return self.shared_embedding(token_ids)
        return self.encoder_embedding(token_ids)

    def decode_input(self, token_ids):
        """Embed decoder input tokens."""
        if self.tie_all:
            return self.shared_embedding(token_ids)
        return self.decoder_embedding(token_ids)

    def project_output(self, hidden_states):
        """Project decoder hidden states to vocabulary."""
        if self.tie_all:
            return torch.matmul(hidden_states, self.shared_embedding.weight.T)
        return self.output_projection(hidden_states)


# Compare parameter counts
enc_dec_tied = EncoderDecoderEmbeddings(vocab_size, hidden_dim, tie_all=True)
enc_dec_untied = EncoderDecoderEmbeddings(vocab_size, hidden_dim, tie_all=False)

Out[21]:

Console

Encoder-Decoder Embedding Parameters:
  Fully tied:     38,400,000 (1 matrix)
  Untied:        153,600,000 (3 matrices)
  Savings:       115,200,000 (75%)

Full tying eliminates 115.2M parameters from the embedding layers alone.

For encoder-decoder models, full weight tying eliminates two-thirds of embedding parameters. T5, one of the most successful encoder-decoder transformers, uses this three-way tying by default. The original paper reported that tying all three matrices worked better than partial tying or no tying for the text-to-text tasks T5 was designed for.

The project_output method still uses torch.matmul(hidden_states, self.shared_embedding.weight.T) even though the "decoder" might have a different hidden dimension scale than the encoder. In practice, T5 uses the same hidden dimension for encoder and decoder, so the shared matrix works without an additional projection. If the encoder and decoder had different hidden dimensions, full tying would not be possible without additional projections.