Weight Tying › Key Parameters

When implementing weight tying in your models, these are the key configuration choices:

• tie_weights (bool): Whether to share the embedding matrix between input and output. Set to True for most language models; set to False when input and output vocabularies differ or when fine-tuning models trained without tying.

• vocab_size (int): Size of the vocabulary. Larger vocabularies increase the parameter savings from tying. With 50K tokens, embedding matrices can dominate smaller models. With 128K tokens, the embedding fraction rises significantly even for medium-depth models.

• hidden_dim or d_model (int): The embedding and hidden dimension. This determines both the embedding matrix size ($V \times d$) and the scale factor ($\sqrt{d}$) used for input scaling.

• embedding_scale (float): Scaling factor applied to input embeddings, typically $\sqrt{d}$. Apply this at forward time rather than modifying the weight matrix, in order to preserve unscaled embeddings for the output projection path.

• bias (bool): Whether to include a bias term in the output projection. Most implementations set this to False when using tied weights. A bias vector would add $V$ parameters and would not be shared, partially defeating the purpose of tying.

• tie_encoder_decoder (bool, encoder-decoder models only): Whether to also tie the encoder input embedding with the decoder embeddings. Recommended when source and target vocabularies are identical, as in T5 and other monolingual sequence-to-sequence models.