如何绑定torch.nn.Embedding与EmbeddingBag的嵌入权重？

需要在神经网络中绑定两个嵌入层的权重：一个是处理顺序敏感Token嵌入的torch.nn.Embedding，另一个是处理顺序无关Token嵌入的torch.nn.EmbeddingBag。手动实现类EmbeddingBag结构时遇到数值稳定性问题，因此希望使用官方的torch.nn.EmbeddingBag，但尝试的权重绑定方法无法正常工作——反向传播后两个层的梯度不一致，代码如下：

#!/usr/bin/env python3

import torch
import torch.nn as nn

if __name__ == "__main__":
    V, max_seq, padding_idx, emb_dim, B = 10, 100, 1, 512, 32

    # 创建嵌入层并初始化EmbeddingBag
    emb_layer = nn.Embedding(V, emb_dim, padding_idx=padding_idx)
    emb_bag = nn.EmbeddingBag.from_pretrained(emb_layer.weight, freeze=False, padding_idx=padding_idx)

    tokens = torch.randint(0, V, (B, max_seq))
    y = torch.randn((B, emb_dim))
    loss = nn.MSELoss()

    # 通过EmbeddingBag反向传播
    y_ = emb_bag(tokens)
    l = loss(y_, y)
    assert emb_bag.weight.grad is None
    l.backward()
    assert emb_bag.weight.grad is not None

    # 断言失败：emb_layer的梯度为空，且两者梯度不一致
    assert emb_layer.weight.grad is not None and \
        torch.allclose(emb_bag.weight.grad, emb_layer.weight.grad)

问题核心是EmbeddingBag.from_pretrained会创建新的参数张量，而非直接共享emb_layer.weight的内存空间，因此反向传播的梯度只会更新emb_bag的权重，无法同步到emb_layer。以下两种方法可实现真正的权重绑定：

方法1：直接赋值weight属性 #

让emb_bag的weight属性直接指向emb_layer的weight张量，实现内存共享：

import torch
import torch.nn as nn

if __name__ == "__main__":
    V, max_seq, padding_idx, emb_dim, B = 10, 100, 1, 512, 32

    emb_layer = nn.Embedding(V, emb_dim, padding_idx=padding_idx)
    # 初始化EmbeddingBag时不使用from_pretrained
    emb_bag = nn.EmbeddingBag(V, emb_dim, padding_idx=padding_idx)
    # 绑定权重
    emb_bag.weight = emb_layer.weight

    tokens = torch.randint(0, V, (B, max_seq))
    y = torch.randn((B, emb_dim))
    loss = nn.MSELoss()

    y_ = emb_bag(tokens)
    l = loss(y_, y)
    l.backward()

    # 断言会成功
    assert emb_layer.weight.grad is not None
    assert torch.allclose(emb_bag.weight.grad, emb_layer.weight.grad)

方法2：使用register_parameter共享参数 #

通过register_parameter方法将emb_bag的weight注册为emb_layer.weight的引用，确保参数共享：

import torch
import torch.nn as nn

if __name__ == "__main__":
    V, max_seq, padding_idx, emb_dim, B = 10, 100, 1, 512, 32

    emb_layer = nn.Embedding(V, emb_dim, padding_idx=padding_idx)
    emb_bag = nn.EmbeddingBag(V, emb_dim, padding_idx=padding_idx)
    # 注册共享参数
    emb_bag.register_parameter('weight', emb_layer.weight)

    tokens = torch.randint(0, V, (B, max_seq))
    y = torch.randn((B, emb_dim))
    loss = nn.MSELoss()

    y_ = emb_bag(tokens)
    l = loss(y_, y)
    l.backward()

    assert emb_layer.weight.grad is not None
    assert torch.allclose(emb_bag.weight.grad, emb_layer.weight.grad)

两种方法都让emb_layer和emb_bag的weight指向同一个张量对象，反向传播时梯度会直接更新这个共享张量，因此两个层的梯度完全一致，实现了真正的权重绑定。

内容的提问来源于stack exchange，提问作者erip