Unexpected Behavior with Weight Sharing between nn.Linear and nn.Embedding

samlk: <h1><a name="p-458199-unexpected-behavior-with-weight-sharing-between-nnlinear-and-nnembedding-in-gpt-2-implementation-1" class="anchor" href="#p-458199-unexpected-behavior-with-weight-sharing-between-nnlinear-and-nnembedding-in-gpt-2-implementation-1"></a>Unexpected Behavior with Weight Sharing between nn.Linear and nn.Embedding in GPT-2 Implementation</h1>
<h2><a name="p-458199-problem-description-2" class="anchor" href="#p-458199-problem-description-2"></a>Problem Description</h2>
<p>I’ve encountered an unexpected behavior when sharing weights between <code>nn.Linear</code> and <code>nn.Embedding</code> layers in my GPT-2 implementation. The order of weight sharing assignment affects the model’s sampling behavior, even with random initialization.</p>
<h2><a name="p-458199-code-3" class="anchor" href="#p-458199-code-3"></a>Code</h2>
<pre data-code-wrap="python"><code class="lang-python">class MLP(nn.Module):
    def __init__(self, d_model: int, activation: Literal['gelu', 'relu'], hidden_dim, **kwargs):
        super().__init__(**kwargs)

        self.d_model = d_model
        self.hidden_dim = hidden_dim

        self.in_layer = nn.LazyLinear(hidden_dim, **kwargs)
        self.activation = nn.GELU(approximate='tanh') if activation == "gelu" else nn.ReLU()
        self.out_layer = nn.LazyLinear(d_model, **kwargs)

    def forward(self, x: torch.Tensor):
        """
        Args:
        :param x: Tensor, shape [batch_size, seq_len, d_model]
        :return: Tensor, shape [batch_size, seq_len, d_model]
        """
        x = self.in_layer(x)
        x = self.out_layer(self.activation(x))
        return x


class MultiHeadedAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int, **kwargs):
        super().__init__(**kwargs)
        assert d_model % num_heads == 0
        self.num_heads = num_heads
        self.d_model = d_model
        self.head_dim = d_model // num_heads

        self.qkv_projection = nn.LazyLinear(3 * d_model, **kwargs)
        self.out_projection = nn.LazyLinear(d_model, **kwargs)

    def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
        B, S, E = x.size()

        query, key, value = self.qkv_projection(x).chunk(3, dim=-1)
        query = query.view(B, S, self.num_heads, E // self.num_heads).transpose(1, 2)
        key = key.view(B, S, self.num_heads, E // self.num_heads).transpose(1, 2)
        value = value.view(B, S, self.num_heads, E // self.num_heads).transpose(1, 2)

       
        out = F.scaled_dot_product_attention(query, key, value, is_causal=True)
        out = out.transpose(1, 2).contiguous().view(B, S, E)
        out = self.out_projection(out)

        return out


class Block(nn.Module):
    def __init__(self, config, **kwargs):
        super().__init__(**kwargs)

        self.d_model = config.d_model
        self.num_heads = config.num_heads
        self.activation = config.activation
        self.hidden_dim = config.hidden_dim

        self.layer_norm1 = nn.LayerNorm(self.d_model, **kwargs)
        self.attention = MultiHeadedAttention(self.d_model, self.num_heads, **kwargs)
        self.layer_norm2 = nn.LayerNorm(self.d_model, **kwargs)
        self.mlp = MLP(self.d_model, self.activation, self.hidden_dim, **kwargs)

    def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
        """
        :param x: torch.Tensor [batch_size, seq_len, d_model]
        :return: torch.Tensor [batch_size, seq_len, d_model]
        """
        x = x + self.attention(self.layer_norm1(x))
        x = x + self.mlp(self.layer_norm2(x))
        return x


class GPT2(nn.Module):
    def __init__(self, config: GPTConfig, **kwargs):
        super().__init__(**kwargs)
        self.config = config
        self.model = nn.ModuleDict(dict(
            tok_embedding=nn.Embedding(config.vocab_size, config.d_model, **kwargs),
            pos_embedding=nn.Embedding(config.max_seq_len, config.d_model, **kwargs),
            blocks=nn.ModuleList([Block(config, **kwargs) for _ in range(config.num_layers)]),
            final_norm=nn.LayerNorm(config.d_model, **kwargs),
        ))
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, **kwargs, bias=False)  # as in gpt2
        # Version A (produces degenerate outputs):
        self.lm_head.weight = self.model.tok_embedding.weight

        # Version B (works correctly)
        self.model.tok_embedding.weight = self.lm_head.weight

    def forward(self, idx: torch.Tensor, targets: torch.Tensor = None):
        """
        :param idx: torch.Tensor [batch_size, seq_len]
        :param targets torch.Tensor [batch_size, seq_len]
        :return: torch.Tensor [batch_size, seq_len, d_model]
        """
        B, T = idx.size()
        pos = torch.arange(0, T, device=idx.device, dtype=torch.long)
        x = self.model.pos_embedding(pos) + self.model.tok_embedding(idx)
        for block in self.model.blocks:
            x = block(x)

        x = self.model.final_norm(x)
        logits = self.lm_head(x)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss

    ```

## Inference Code and Observed Behavior
```python
model = GPT2(GPTConfig())
model.eval()
model.to(device)

max_length = 30
seq_to_sample = 5

tokenizer = tiktoken.get_encoding('gpt2')
tokens = tokenizer.encode("Hello, I'm a language model,")
tokens = torch.tensor(tokens, dtype=torch.long, device=device)
tokens = tokens.unsqueeze(0).repeat(seq_to_sample, 1)

outputs = model.sampling_loop(tokens, max_length)

for i in range(seq_to_sample):
    out_tokens = outputs[i, :max_length].tolist()
    decoded = tokenizer.decode(out_tokens)
    print(f"&gt;{decoded}")
</code></pre>
<h3><a name="p-458199-output-with-version-a-buggy-4" class="anchor" href="#p-458199-output-with-version-a-buggy-4"></a>Output with Version A (buggy)</h3>
<pre><code class="lang-auto">&gt;Hello, I'm a language model,,,,,,,,,,,,,,,,,,,,,,,
&gt;Hello, I'm a language model,,,,,,,,,,,,,,,,,,,,,,,
&gt;Hello, I'm a language model,,,,,,,,,,,,,,,,,,,,,,,
&gt;Hello, I'm a language model,,,,,,,,,,,,,,,,,,,,,,,
&gt;Hello, I'm a language model,,,,,,,,,,,,,,,,,,,,,,,
</code></pre>
<p>The model with Version A (where <code>self.lm_head.weight = self.model.tok_embedding.weight</code>) produces degenerate outputs, repeatedly sampling the same token (comma). However, when using Version B (<code>self.model.tok_embedding.weight = self.lm_head.weight</code>), the model produces more diverse and sensible continuations.</p>
<p>Why does the order of weight sharing assignment matter in this case, even with random initialization?</p>

ptrblck: <p>Your model is missing the <code>forward</code> function and should thus directly fail with:</p>
<pre data-code-wrap="python"><code class="lang-python">    raise NotImplementedError(

NotImplementedError: Module [GPT2] is missing the required "forward" function
</code></pre>
<p>since <code>self(x)</code> will call into <code>self._call_impl(*args, **kwargs)</code> and <code>forward_call(*args, **kwargs)</code>.</p>

samlk: <p>I omitted it, because it’s long. Now it’s there.</p>