In Section 3.4 of "Attention Is All You Need" (Vaswani et al., 2017), the authors describe sharing weights between two embedding layers and the pre-softmax linear transformation, then multiplying the embedding weights by the square root of the model dimension. What is the primary purpose of this weight-sharing scheme?
- It lets the encoder and decoder share the same set of self-attention weights, so a single set of query/key/value matrices is reused in every layer of the network
- It reduces the depth of the network by merging the final feed-forward sublayer directly into the output embedding matrix
- It ties the same learned matrix to both converting tokens into vectors and converting the decoder's final vectors back into vocabulary logits, reducing the total parameter count and linking the input and output token representations
- It guarantees that every possible output token receives an identical probability at initialization, which is later corrected during fine-tuning
Why C? And why not the others?
Correct answer: C. It ties the same learned matrix to both converting tokens into vectors and converting the decoder's final vectors back into vocabulary logits, reducing the total parameter count and linking the input and output token representations
Vaswani et al. tie a single learned weight matrix to both roles: converting input and output tokens into d_model-dimensional vectors and, transposed, converting the decoder's final vectors into logits over the vocabulary before the softmax; because one matrix serves both purposes, the model needs fewer total parameters than if the embedding and output projection were learned separately, and the input and output token representations remain linked to each other, an idea the paper credits to prior work on tying input and output embeddings. The embedding weights are then multiplied by the square root of the model dimension to keep their scale comparable to the positional encodings added to them. The option describing shared encoder/decoder self-attention weights is wrong because the tied weight matrix here is the token embedding/output-projection matrix, not the query, key, or value projections used inside attention, which remain separate learned parameters in each layer. The option describing merging the feed-forward sublayer into the embedding matrix is wrong because the position-wise feed-forward network is a separate sublayer with its own weights that this tying scheme does not touch. The option claiming every output token receives an identical initial probability is wrong because the tied matrix is initialized with the same random values used for embeddings, not with values that force a uniform output distribution.
Source: Vaswani et al., "Attention Is All You Need" (2017), arXiv:1706.03762, Section 3.4 (Embeddings and Softmax)