Weight Tying in Transformer Embeddings

Language models contain two large embedding matrices that seem to serve completely different purposes. One converts input tokens into continuous vectors at the start of the network. The other converts output vectors back into token probabilities at the end. If you examine these matrices for a model like GPT-2, you will find that each has 50,257 rows, one per vocabulary token, and 768 columns, one per hidden dimension. Together, they account for roughly 77 million parameters in a model whose total size is only about 117 million parameters. That means over 65 percent of all parameters in GPT-2 Small live in just these two matrices.

What makes this extraordinary is that both matrices answer fundamentally the same question: what does each vocabulary token mean within this model's learned representation space? The input embedding asks, "given this token, what vector should I pass into the transformer?" The output projection asks, "given this hidden state, how strongly does it correspond to each token?" These are two sides of the same coin. A model that has learned a rich representation for the word "justice" should encode it consistently whether reading it or generating it.

Weight tying exploits this insight with striking simplicity: instead of learning two separate matrices, the model uses one. The same embedding matrix that converts "justice" into a vector at the input also is the output target when the model wants to generate "justice". This single constraint halves the embedding-related parameters and, counter-intuitively, often improves model quality. You might expect that giving the model fewer degrees of freedom would hurt performance, but the opposite tends to be true. The constraint acts as a useful regularizer that forces coherent representations.

This chapter covers weight tying from its mathematical foundations to its practical implementation. We will understand exactly why tying makes theoretical sense, trace through numerical examples to see how shared embeddings work in practice, examine the gradient dynamics that make tied models learn differently, and explore encoder-decoder extensions where three matrices can be collapsed into one. We will also study the limitations: tying is not always appropriate, and knowing when to avoid it is as important as knowing how to implement it.

By the end of this chapter, you will understand one of the most universally adopted architectural choices in modern language modeling, appearing in GPT-2, BERT, T5, and nearly every transformer trained since 2017.

Every language model that processes text through a transformer architecture faces the same fundamental challenge at its boundaries. Transformers operate on continuous vectors, but text consists of discrete tokens. Bridging this gap requires learned mappings in both directions: from tokens to vectors at the input, and from vectors back to tokens at the output. These two mappings are implemented as the input embedding matrix and the output projection matrix, and understanding them precisely is the first step toward appreciating why tying them together makes sense.

Think of the transformer as a black box that maps sequences of dense vectors into new sequences of dense vectors. Before the transformer can do its work, you need to convert your token indices into vectors it can process. After the transformer produces its output, you need to convert those output vectors back into probabilities over the vocabulary. The two embedding matrices handle these conversion tasks.

Both matrices have the same shape, both operate over the same vocabulary, and both are involved at every single forward pass through the model. Yet in naive implementations, they are treated as entirely separate learned objects. Weight tying questions whether this separation is necessary or even desirable.

The input embedding matrix $\mathbf{E} \in \mathbb{R}^{V \times d}$ is one of the most fundamental components in any neural language model. It stores a learned dense vector representation for each token in the vocabulary. When the model sees token index $t$, it retrieves the corresponding row from this matrix:

where:

• $\mathbf{E}$: the input embedding matrix of shape $V \times d$, where every row contains the learned representation of one vocabulary token
• $t$: the input token index, an integer from 0 to $V - 1$
• $\mathbf{h}_{\text{in}}$: the resulting embedding vector of dimension $d$, which enters the transformer's processing pipeline
• $V$: the vocabulary size (50,257 for GPT-2, 32,000 for LLaMA-2, 128,000 for newer models)
• $d$: the embedding dimension, also called the model dimension or hidden dimension

This operation is purely a table lookup. Token 42 retrieves row 42 from the matrix. There is no computation beyond the indexing, which is one reason why embedding lookups are extremely fast in practice. The entire operation can be implemented as a gather operation on the weight matrix.

Over the course of training, the rows of $\mathbf{E}$ are shaped by gradient descent to capture the semantic and syntactic properties of each token. Tokens that appear in similar contexts will develop similar embeddings. Tokens that serve similar grammatical roles will develop representations that are geometrically related. The embedding matrix effectively becomes a learned dictionary that maps the discrete world of tokens to the continuous world of vector spaces.

The output projection matrix $\mathbf{W}_{\text{out}} \in \mathbb{R}^{V \times d}$, also called the language model head or LM head, performs the inverse operation. After the transformer has processed the input sequence and produced a final hidden state $\mathbf{h}_{\text{out}} \in \mathbb{R}^d$, the model needs to decide which token should come next. This decision is made by computing a score for every token in the vocabulary:

where:

• $\mathbf{z}$: the output logits vector of length $V$, one unnormalized score per vocabulary token
• $\mathbf{W}_{\text{out}}$: the output projection matrix of shape $V \times d$, where each row encodes the "target pattern" for one vocabulary token
• $\mathbf{h}_{\text{out}}$: the final hidden state from the last transformer layer, a vector of dimension $d$

Each row of $\mathbf{W}_{\text{out}}$ contains what you might call the "output embedding" for one vocabulary token. The full matrix multiplication computes the dot product between the hidden state and each token's output embedding in one vectorized operation, yielding a score for every token in the vocabulary simultaneously.

After computing these logits, we apply the softmax function to convert them into a proper probability distribution over the vocabulary for the next token prediction:

where:

• $P(\text{token}_i \mid \text{context})$: the probability of token $i$ being the correct next token given the context
• $z_i$: the logit for token $i$, computed as $\mathbf{w}_i \cdot \mathbf{h}_{\text{out}}$ where $\mathbf{w}_i$ is row $i$ of $\mathbf{W}_{\text{out}}$
• $\exp(z_i)$: the exponential applied to logit $i$, which ensures all values are positive
• $\sum_{k=1}^{V} \exp(z_k)$: the sum of all exponentials across the full vocabulary, serving as a normalizing constant

Why does this formula make sense? Notice that taking the exponential of each logit guarantees positive values regardless of the sign of the original logit, which is a prerequisite for valid probabilities. The denominator then ensures all values sum to exactly 1. Tokens with higher logits, those whose output embeddings more closely align with the hidden state, receive higher probabilities.

Notice the dimensional symmetry: both $\mathbf{E}$ and $\mathbf{W}_{\text{out}}$ have exactly the same shape, $V \times d$. Both have $V$ rows, one per vocabulary token. Both have $d$ columns, one per hidden dimension. Both map between the same two spaces: the discrete token space with $V$ elements and the continuous embedding space with $d$ dimensions.

This symmetry is not a coincidence. It reflects the deeper fact that both matrices perform the same abstract function: they define a correspondence between tokens and points in a $d$-dimensional vector space. The input embedding says "token $t$ lives at point $\mathbf{e}_t$." The output projection says "a hidden state should produce token $t$ with probability proportional to how much it resembles $\mathbf{w}_t$." These two statements describe the same geometric relationship from opposite directions.

With a vocabulary of 50,000 tokens and hidden dimension of 768, each matrix contains over 38 million parameters. Together, they account for nearly 77 million parameters. For a model like GPT-2 Small, which has about 117 million total parameters, this means over 65 percent of the model's capacity lives in these two identically sized matrices. The question weight tying asks is: do we really need both?

We have established that language models maintain two large matrices with identical shapes: one for reading tokens and one for generating them. Both have dimensions $V \times d$, where each row represents one vocabulary token. This structural symmetry hints at a deeper connection worth pursuing carefully.

Think about what information these matrices encode. The input embedding learns "what does this token mean when I read it?" The output projection learns "what hidden state pattern should produce this token?" But at a deeper level, these are really two perspectives on the same question: what is the semantic identity of this token within the model's learned representation space? A token's meaning should not depend on whether you are reading it or generating it.

Think of it this way. If you learned the word "melancholy" in school, you would use the same mental concept whether you are recognizing that word in a sentence you are reading or deciding to write it in a sentence you are composing. You do not maintain one neural representation of "melancholy" for reading and a completely separate representation for writing. You have one underlying concept, and it serves both purposes. Weight tying imposes the same discipline on language models.

a model's representation of any token should be internally consistent: the vector that captures "what this token means as input" should be the same vector that is "the target to match when generating this token as output." Any discrepancy between these two representations is wasted degrees of freedom at best, and incoherence at worst.

Weight tying formalizes this intuition by collapsing both matrices into one:

where:

• $\mathbf{W}_{\text{out}}$: the output projection matrix, now set equal to $\mathbf{E}$ rather than being an independent learned parameter
• $\mathbf{E}$: the input embedding matrix of shape $V \times d$, which now serves both roles simultaneously

This single equation halves the embedding-related parameters. But what does it mean computationally? Let's trace through the math to understand how the shared matrix serves both roles in a coherent way.

When a token enters the model, we look up its embedding. This operation is unchanged by tying:

The token at index $t$ retrieves row $t$ from the matrix, a simple table lookup. Training updates this row to encode useful information about token $t$ for downstream processing.

When the transformer produces its final hidden state $\mathbf{h}_{\text{out}}$, we convert this $d$-dimensional vector into scores for all $V$ vocabulary tokens. With weight tying, we use the same matrix for this conversion:

This matrix multiplication produces a vector of $V$ logits. But what is happening geometrically inside this operation? Each logit is a dot product between the hidden state and one row of the embedding matrix.

The logit for token $i$ decomposes as follows. We need a formula that takes the hidden state $\mathbf{h}_{\text{out}}$ and produces a score reflecting how likely token $i$ is to be the correct next word. With tied weights, we compute:

where:

• $z_i$: the logit (unnormalized score) for token $i$, a single real number
• $\mathbf{e}_i$: the embedding vector for token $i$ (row $i$ of $\mathbf{E}$), a $d$-dimensional vector with components $e_{i,1}, e_{i,2}, \ldots, e_{i,d}$
• $\mathbf{h}_{\text{out}}$: the final hidden state vector, with $d$ components $h_1, h_2, \ldots, h_d$
• $d$: the embedding dimension

Why does this formula make sense? Notice that the dot product measures geometric alignment between two vectors. Two vectors pointing in the same direction in the $d$-dimensional embedding space yield a large positive value. Orthogonal vectors yield zero. Vectors pointing in opposite directions yield a large negative value. Softmax then converts these alignment scores into probabilities, so the token whose embedding most closely aligns with the hidden state receives the highest probability.

This creates a beautiful geometric interpretation: the transformer's job is to produce a hidden state that points in the direction of the correct next token's embedding. When processing the context "The cat sat on the ___", the transformer should generate a hidden state that is close to the embedding vectors for words like "mat," "floor," or "couch," and far from the embedding vectors for words like "quantum" or "bureaucracy."

With tied weights, the embedding vectors that serve as targets for this similarity search are exactly the same vectors used to represent tokens in the input. The model must learn one coherent embedding for each token that works both as an input representation and as an output target. This constraint is not a restriction; it is a clarification that forces the model to build consistent semantics.

Translating this math into code is straightforward. We create a single embedding matrix and use it for both encoding (lookup) and decoding (matrix multiplication):

The parameter savings are substantial. We eliminated one entire $V \times d$ matrix, cutting embedding-related parameters in half. For models with large vocabularies or smaller hidden dimensions, this reduction represents a significant fraction of total model size. More importantly, as we will see shortly, this reduction often comes with improved performance, not degraded performance.

Weight tying reduces memory use, but its theoretical justification also explains why input and output embeddings should be related. Understanding this justification deepens your intuition for how transformers represent language.

Consider what each embedding encodes over the course of training. The input embedding for token $t$ gets updated whenever token $t$ appears in the input, with gradients flowing backward from the loss through all the transformer layers above. It learns a representation that makes downstream processing easier: similar tokens develop similar embeddings, tokens that play similar grammatical roles cluster together, and the geometric structure of the embedding space reflects the distributional statistics of the language.

The output embedding for token $t$ gets updated whenever token $t$ should be predicted as the next word, with gradients flowing directly from the cross-entropy loss. It learns what kind of hidden state the transformer should produce when "token $t$ is the right answer." The hidden state that generates "cat" should look similar to the output embedding for "cat."

Both representations capture the same fundamental thing: the contextual signature of token $t$ within the model's learned world model. A token's input representation and its output representation are both shaped by the same distributional statistics of the language. They are really two approximations to the same underlying concept, learned through slightly different optimization paths.

Here is a simple argument for why tying should help. Suppose you have two separate matrices, $\mathbf{E}$ and $\mathbf{W}_{\text{out}}$, and the model learns them independently. After training, you look at the embedding for "dog" in the input matrix and the embedding for "dog" in the output matrix. If the model has learned well, these two vectors should be similar, because both capture the meaning of "dog." But they are not required to be similar, so the model wastes capacity representing the same concept twice in slightly different ways.

Weight tying removes this redundancy by decree. The model is forced to find a single representation that works for both purposes, which means every gradient update contributes to a more coherent and consistent token representation. You are effectively doubling the amount of learning signal that shapes each token's embedding, because both the input path and the output path update the same weights.

This consistency argument applies directly to language modeling, where the input and output vocabularies are identical. The model reads and writes in the same language, using the same tokens. There is no fundamental reason why the concept of "dog" should be represented differently depending on whether the model is reading it or generating it.

Research has consistently confirmed that weight tying helps rather than hurts performance:

• Press and Wolf (2017) introduced the technique systematically and demonstrated that tying input and output embeddings improves perplexity across language modeling benchmarks, despite having fewer parameters. This was the result that established weight tying as a standard practice.

• Inan et al. (2017) provided complementary theoretical analysis, connecting weight tying to a regularization perspective where the shared matrix is encouraged to satisfy both input and output objectives simultaneously.

• Practical adoption: Most modern language models, including GPT-2, BERT, RoBERTa, T5, and their descendants, use weight tying by default. The technique has become so standard that its use is often not even mentioned in architecture descriptions.

The improvement in perplexity also reflects more than the regularization from having fewer parameters. Tied embeddings tend to learn better representations because gradients from the output loss flow directly into the input embeddings, and gradients from the input embedding usage flow into the output predictions. The two optimization objectives reinforce each other.

Let's implement a complete language model head with weight tying, handling the practical details that matter in real systems. A production-grade implementation needs to manage positional encodings, dropout, and the subtleties of the embedding scale.

The key engineering insight in a full implementation is that the embedding matrix is shared by reference, not by copying. Both the encode and decode operations point to the same underlying tensor. When gradients update the shared matrix, they simultaneously improve both the input representations and the output targets. This is not a coincidence or a side effect. It is the core mechanism that makes weight tying beneficial.

The key implementation detail is in the project method. With tied weights, we directly use self.token_embedding.weight.T instead of a separate projection layer. This ensures that gradients flow through the same parameters during both forward and backward passes. There is no additional memory allocated for a second embedding matrix. The model literally shares the same tensor object.

Notice that position embeddings are not tied and do not participate in the output projection. This is by design. Position embeddings encode sequential position information that has no meaning in the output vocabulary space. Only the token embeddings, which capture the semantic identity of each vocabulary item, are shared between input and output.

Let's verify that our implementation shares parameters correctly and that gradients flow as expected:

The gradient from the output loss flows directly into the token embedding matrix, confirming that the tying works correctly. With tied weights, every backward pass accumulates gradients from multiple sources into the same underlying parameters. This is the mechanism by which tied models receive more learning signal per parameter.

The formulas above describe weight tying abstractly, but seeing the mechanism in action makes it concrete. Let's build a tiny language model and trace through exactly how tied embeddings convert a hidden state into token probabilities. We will work with small enough numbers to inspect every value manually.

We will work with a vocabulary of just seven words using 4-dimensional embeddings. These small numbers let us inspect every value and understand exactly what is happening at each step. In a real model you might have 50,000 words and 768 dimensions, making manual inspection impractical, but the same mathematical operations apply at every scale.

Each word has a 4-dimensional embedding vector. These are randomly initialized, as they would be before training. In a trained model, semantically similar words would have similar embeddings. The words "cat" and "dog" would cluster together because they are both animals that appear in similar grammatical contexts. The words "sat" and "ran" would form another cluster as verbs. Weight tying means that these same clusters serve as targets for the output prediction: when the model wants to generate "cat," it aims to produce a hidden state that points in the direction of the "cat" embedding vector.

Now comes the key insight. Suppose the transformer has processed some context and produced a final hidden state. With weight tying, we predict the next token by computing dot products between this hidden state and every embedding in our vocabulary.

Let's simulate this. We'll create a hidden state that is similar to the "cat" embedding, as if the model were about to predict "cat" as the next word. This simulates what would happen if the transformer correctly learned to associate some context with the concept of "cat":

The hidden state is intentionally constructed to align with "cat"'s embedding while moving away from the average of the other token embeddings. In a real model, the transformer layers would produce this hidden state based on the input context. The contrastive offset represents the fact that the model's output is not a perfect copy of any single embedding; it is a point in the embedding space that aligns more strongly with some tokens than others.

Now let's examine what happens inside the forward pass. For each vocabulary token, we compute the dot product between the hidden state and that token's embedding. This gives us a raw score, which softmax then converts into a probability:

The dot product reveals the alignment between the hidden state and each embedding. "Cat" receives the highest score because its embedding is most similar to the hidden state we constructed. After softmax normalization, this translates to the highest probability. The other tokens receive lower but non-zero probabilities. This reflects the fact that uncertainty exists about the correct next token.

This is the heart of weight tying: the same embedding that represents "cat" for input processing is also the target pattern the model aims to produce when generating "cat". The transformer's job is to transform the input context into a hidden state that matches the correct next token's embedding. Learning happens both by improving the embedding representations and by improving the transformer's ability to produce hidden states that match those representations.

Let us trace the complete computation from logits to probabilities for the first three tokens to make the math fully concrete. Suppose the hidden state produces logits of $z = [0.8, 1.4, 0.3, 0.1, 0.5, 0.2, 0.4]$ for our seven tokens. The softmax computation proceeds:

The normalizing denominator sums all seven exponentials, and each token's probability is its exponential divided by this sum. The token with the highest logit, "cat" in this example, receives the highest probability. The key observation is that with tied weights, these logit values are dot products with the input embedding vectors, so improving the input embedding for "cat" simultaneously makes it easier to predict "cat" as an output.

To better understand how the dot product works as a similarity measure, let's visualize the pairwise similarities between all embeddings in our tiny vocabulary along with the hidden state:

The heatmap reveals the geometric structure of our embedding space. The bottom row and rightmost column show how similar the hidden state is to each vocabulary embedding. Since we constructed the hidden state to resemble "cat," that cell shows the highest similarity in the hidden row. The diagonal entries reflect each token's self-similarity. Off-diagonal entries show how similar different tokens are to each other in the learned representation space.

With tied weights, this same similarity matrix governs both how tokens are processed as input and how they are selected as output. Making "cat" more recognizable as an input (by moving its embedding to a more distinctive location in the space) simultaneously makes it easier to generate "cat" as an output (by making the target point more distinct from other tokens).

Weight tying becomes more or less impactful depending on how large the vocabulary is relative to the rest of the model. In a tiny model, the embedding matrices might account for the majority of parameters. In a massive model with many deep transformer layers, the embedding fraction shrinks. Understanding this scaling relationship helps you calibrate how much you should care about tying in different contexts.

The fundamental relationship is simple: the number of embedding parameters grows linearly with vocabulary size and embedding dimension, while the number of transformer layer parameters grows quadratically with the hidden dimension (due to the $4d^2$ parameters in attention and feedforward components). As models scale up in depth, the relative importance of the embeddings decreases.

The "Large vocab" configuration shows an interesting pattern: models with bigger vocabularies, like those using 128K-token vocabularies for multilingual support or code generation, benefit more from weight tying because the embedding matrix is a larger fraction of total parameters. As newer models adopt increasingly large vocabularies to support diverse scripts and domains, weight tying becomes progressively more important as an efficiency measure.

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:

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.

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.

Weight tying does not just reduce parameters. It changes the fundamental dynamics of how the model learns, and understanding these dynamics helps explain why tied models often outperform untied models with the same number of free parameters.

The most important change is in how gradients flow. With separate embedding matrices, the input embedding receives gradients only through the transformer layers, while the output projection receives gradients directly from the cross-entropy loss. With tied weights, the single shared matrix receives gradients from both sources simultaneously. This is a form of implicit multi-task learning: the embedding must simultaneously satisfy two objectives, and each objective provides additional signal about what good representations look like.

Consider a rare token like a technical term that appears infrequently in the training corpus. With separate embeddings, the input embedding for this token receives sparse gradient updates (only when it appears in the input), and the output embedding receives sparse updates (only when it should be predicted). With tied weights, both sources of gradient contribute to the same parameters. This provides more learning signal per occurrence. This is particularly useful for vocabulary items that are infrequent as either input or output but not both.

With tied weights, the embedding matrix receives gradients from two sources:

1. Input gradients: Backpropagated through the transformer from the loss. When a token appears in the input, its embedding affects all downstream computations, and gradients flow backward through all the transformer layers to update the embedding.

2. Output gradients: Direct gradients from the output projection. When a token should be predicted, the dot product between the hidden state and that token's embedding directly contributes to the cross-entropy loss. This provides a gradient that says "move this embedding closer to the current hidden state."

This double gradient flow can be viewed as implicit multi-task learning. The embedding must simultaneously satisfy two objectives: representing tokens well for input processing so the transformer can reason about them, and providing good targets for output prediction so the model can accurately identify which token to generate next.

The gradient magnitude shows how much the embedding weights would change in a single training step (before learning rate scaling). With tied weights, this gradient is typically larger than it would be with separate embeddings because it aggregates signals from both the input and output paths. This can speed up learning for rare tokens that might otherwise receive sparse gradient updates, since each appearance provides double the learning signal.

The combined gradient also has a regularizing effect. The input gradient pushes the embedding toward representations that make downstream processing easier. The output gradient pushes it toward representations that align well with the hidden states the transformer produces for contexts where this token is appropriate. When both gradients point in compatible directions, the embedding is in a region of representation space that satisfies both constraints simultaneously, which is exactly where you want it.

A subtle issue arises with weight tying: the optimal scale for input embeddings may differ from the optimal scale for output projections. This affects real training dynamics and model performance.

Input embeddings are often scaled by $\sqrt{d}$ before being added to positional encodings, following the original transformer paper by Vaswani et al. (2017). The scaling formula is:

where:

• $\mathbf{h}_{\text{scaled}}$: the scaled input embedding that enters the transformer
• $d$: the embedding dimension
• $\mathbf{E}[t]$: the raw embedding lookup for token $t$
• $\sqrt{d}$: the scaling factor, which counterbalances the variance reduction that occurs when embeddings are initialized with small standard-deviation values

Without this scaling, embeddings initialized with small weights (standard deviation around 0.02) would be numerically negligible compared to the positional encodings, which are also on a small scale. The positional encodings would dominate the input signal, preventing the model from distinguishing token identities. The $\sqrt{d}$ scaling amplifies the token embeddings to a magnitude that balances both sources of information.

However, if we applied this scaling directly to the weight matrix, it would also affect the output logits when using tied weights. Scaled embeddings as output targets would produce much larger logits than expected, potentially destabilizing softmax and making the output distribution unreasonably peaked.

Modern implementations handle this by applying scaling at input time rather than modifying the embedding matrix itself:

The scale factor of approximately 27.7 for hidden dimension 768 significantly amplifies the input embeddings. By applying this scaling at runtime rather than baking it into the weights, we decouple the input and output requirements. The transformer processes scaled embeddings while the output projection uses raw embeddings, allowing each path to operate at its optimal magnitude. This is an example of the broader principle that tying weights does not mean treating both uses identically; you can still apply path-specific transformations.

Weight tying is not always the right choice. Knowing when to apply it requires understanding both the benefits and the situations where the constraints it imposes cause more harm than good.

The benefits are clearest when input and output vocabularies are identical and the model needs to be efficient. Language modeling is the canonical case: you read text and generate text in the same language using the same tokens. The input and output domains are perfectly symmetric, and there is no principled reason why token representations should differ between these roles.

The constraints hurt most when input and output require different representations. Machine translation is the clearest counterexample: if you are translating from English to French using separate tokenizers, you cannot tie embeddings because the vocabularies are literally different. But even with a shared multilingual vocabulary, tying may be suboptimal for translation because the optimal encoding for reading an English source sentence may differ from the optimal target for generating a French output sentence. The distributional statistics of English and French are different enough that forcing a single shared representation for all tokens may limit what the model can learn.

Fine-tuning scenarios present another consideration. If you are fine-tuning a model that was pre-trained without weight tying, introducing tying during fine-tuning would require merging two separately optimized matrices, which typically hurts performance. The safest approach is to match the architecture of the pre-trained model.

The heuristic reveals an important pattern: smaller models like GPT-2 Small have embeddings that constitute a significant fraction of total parameters (over 15 percent), making weight tying highly impactful. For massive models like GPT-3 175B, embeddings are less than 1 percent of parameters, so tying provides minimal savings. However, even large models typically use weight tying because it rarely hurts performance and provides a small memory benefit along with improved gradient flow. The multilingual translation case shows when tying is impossible: different input and output vocabularies require separate embedding spaces.

When making the tying decision in practice, consider these factors:

• Same vocabulary: Tying is only meaningful when input and output tokens come from the same vocabulary. This is true for language modeling, text classification with output over vocabulary, and sequence generation in the same language.

• Model size: For models where embeddings exceed 10 percent of total parameters, tying provides noticeable savings. For very large models, the savings are proportionally smaller but the qualitative benefits of consistent representations remain.

• Training from scratch: When training from scratch, tying is generally the right default choice. The model will naturally learn representations that satisfy both constraints simultaneously.

• Fine-tuning: Match the pre-trained model's configuration. Changing the tying configuration during fine-tuning disrupts the learned representation structure.

• Domain symmetry: If the input and output domains are symmetric (reading and writing in the same language, generating the same types of content you process as input), tying is appropriate. If there is an asymmetry in what optimal input versus output representations look like, consider whether untied weights might give the model more flexibility.

Weight tying represents one of those elegant techniques where reducing complexity improves results. The constraint that input and output embeddings share the same learned representation forces the model to develop more coherent internal semantics, and the research record shows consistent improvements in perplexity when tying is applied appropriately.

The primary limitation of weight tying is inflexibility. When input and output tasks require different token representations, tied weights create tension that the model cannot resolve without compromising both objectives. Machine translation between languages with different scripts is the canonical example: the optimal encoding for reading Japanese may differ substantially from the optimal target representation for generating English output, even if both pass through the same shared vocabulary of BPE tokens. In such cases, untied weights give the model the freedom to specialize each matrix for its respective task. Forcing a single matrix to serve both roles when the tasks differ is asking the model to make an impossible compromise.

There is also a subtler capacity argument. Very large models may benefit from the additional expressiveness of separate embeddings. When parameter count is less constrained, the regularization effect of tying matters less, and the model might learn better with independent representations that can specialize. However, the empirical evidence here is mixed. GPT-3 at 175 billion parameters uses weight tying, suggesting that even at that scale, the benefits of consistency outweigh the benefits of separate specialized matrices. The general consensus in the field is that tying is appropriate unless you have a specific reason to untie, such as different vocabularies or clear empirical evidence that untied weights perform better for your specific task.

A practical limitation arises in fine-tuning scenarios. If a model was pre-trained without weight tying and you want to apply weight tying during fine-tuning to reduce memory usage, you face the problem of how to initialize the shared matrix from two separately optimized matrices. Simply averaging them or choosing one over the other disrupts the pre-trained representations, often hurting performance. This makes weight tying a design choice that should be made at the beginning of training rather than retrofitted later.

The embedding scale issue discussed earlier in this chapter is also a practical limitation. The $\sqrt{d}$ scaling required for numerical stability during input processing is not appropriate for the output projection path. Implementations must carefully apply scaling at the right point in the forward pass, and this requirement is easy to get wrong when implementing weight tying for the first time. Getting the scaling wrong can cause training instability or suboptimal performance even when the conceptual implementation is correct.

Despite these limitations, weight tying's impact on the field has been substantial. It has become so universally adopted that its presence is often assumed rather than stated in architecture descriptions. The technique has influenced thinking about what neural networks are learning: if input and output embeddings converge toward similar values during training, that is evidence that the model is discovering a unified representation space rather than two separate approximations to the same underlying concept. Weight tying simply makes this convergence explicit and mandatory, saving parameters and accelerating learning in the process.

The broader lesson from weight tying is about recognizing and exploiting symmetry in model architecture. Whenever two components of a model are doing conceptually similar things, tying their weights may be beneficial. This principle extends beyond embeddings to other shared-weight architectures, like Siamese networks for similarity tasks or shared encoder components in multi-task models. Weight tying is a specific trick that illustrates a general principle: structural constraints reflecting meaningful symmetries tend to produce better models.

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.

Weight tying exploits the structural symmetry between input embedding and output projection matrices in language models. Both matrices have shape $V \times d$, where $V$ is vocabulary size and $d$ is the hidden dimension. Instead of maintaining separate matrices $\mathbf{E}$ for input and $\mathbf{W}_{\text{out}}$ for output, weight tying sets $\mathbf{W}_{\text{out}} = \mathbf{E}$, so a single shared matrix serves both roles.

The key insights from this chapter:

• Parameter reduction: Weight tying eliminates one full embedding matrix, saving $V \times d$ parameters. For a 50K vocabulary with 768-dimensional embeddings, that is 38 million fewer parameters, often more than 15 percent of a small model's total capacity.

• Theoretical justification: Input and output embeddings both answer the question "what does this token mean?" Tying them forces a consistent semantic representation where the hidden state that produces a token is similar to that token's embedding. This is not a constraint; it is a clarification of what good representations should look like.

• Dot product as similarity: With tied weights, generating a token corresponds to finding which embedding most closely aligns with the current hidden state. The transformer's task is to produce hidden states that are similar to the embeddings of the correct next tokens.

• Gradient dynamics: Tied weights receive gradients from both the input path (through transformer layers) and the output path (directly from cross-entropy loss). This provides more learning signal per token occurrence and acting as implicit multi-task learning.

• Encoder-decoder extension: In sequence-to-sequence models, all three embedding matrices (encoder input, decoder input, decoder output) can share weights when the vocabulary is shared, eliminating two-thirds of embedding parameters. T5 uses this approach.

• Embedding scale: The $\sqrt{d}$ scaling required for input stability should be applied at forward time, not baked into the weight matrix, so that the output projection sees unscaled embeddings.

• When to use: Weight tying is recommended when vocabularies are large relative to model depth and input and output domains are symmetric. Avoid it when vocabularies differ, when fine-tuning models trained without tying, or when empirical evidence in your specific task suggests separate matrices perform better.

Modern language models nearly universally adopt weight tying, making it a fundamental architectural decision rather than an optimization. Understanding why it works deepens your understanding of what these models are learning: a unified representation space where token identity is consistently encoded whether a token is being read or generated.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about weight tying in language models.