4、权重共享实现方式：手动实现权重共享、使用nn.Parameter共享、跨层共享技巧

权重共享，说白了就是让不同层共用同一份参数。为什么要这么做？我刚开始接触时也觉得多此一举，直到在项目中遇到参数量爆炸的问题——一个深层网络光Embedding层就占了80%的内存。嗯，这时候权重共享的价值就体现出来了。

这一节我们聊三种实现方式。每种都有它的适用场景，我按从「原始」到「优雅」的顺序来讲。

4.1 手动实现权重共享：最笨但最可控

最直接的方式，就是自己维护一个权重矩阵，然后手动传给不同的层。我在早期做词向量对齐项目时用过这招，虽然代码丑了点，但逻辑非常清晰。

import torch
import torch.nn as nn

class ManualSharedEmbedding(nn.Module):
    def __init__(self, vocab_size, embed_dim):
        super().__init__()
        # 自己创建一个权重矩阵
        self.shared_weight = torch.randn(vocab_size, embed_dim)
        # 注册为参数，让优化器能更新它
        self.shared_weight = nn.Parameter(self.shared_weight)
        
    def forward(self, input_ids_a, input_ids_b):
        # 手动查表
        emb_a = nn.functional.embedding(input_ids_a, self.shared_weight)
        emb_b = nn.functional.embedding(input_ids_b, self.shared_weight)
        return emb_a, emb_b

这里有个坑：nn.functional.embedding默认不会把权重注册到模型中。如果你直接用torch.randn创建，梯度根本传不回来。我曾经在这上面浪费了半天——模型训练loss纹丝不动，查了半天才发现权重根本没被优化器追踪。

4.2 使用nn.Parameter共享：更优雅的方式

手动方式虽然可控，但代码量偏大。我个人更推荐用nn.Parameter配合nn.Embedding来做共享。你想想看，PyTorch的nn.Embedding底层就是一个nn.Parameter，我们直接复用它的实例就行。

class ParameterSharedEmbedding(nn.Module):
    def __init__(self, vocab_size, embed_dim):
        super().__init__()
        # 创建一个Embedding层
        self.embed = nn.Embedding(vocab_size, embed_dim)
        
    def forward(self, input_ids_a, input_ids_b):
        # 两次调用用的是同一个embed层
        emb_a = self.embed(input_ids_a)
        emb_b = self.embed(input_ids_b)
        return emb_a, emb_b

这种方式的好处是：nn.Embedding内部已经处理好了参数注册、梯度传播、设备迁移等细节。你只需要保证多个地方引用的是同一个对象就行。

我记得有一次，同事把两个Embedding层分别定义，然后手动复制权重——结果训练时一个在更新，另一个纹丝不动。这就是典型的「伪共享」。说白了，共享的不是参数对象，而是参数值，那梯度当然传不过去。

4.3 跨层共享技巧：不只是Embedding

跨层共享的应用场景更广。比如Transformer中，Encoder和Decoder的Embedding层可以共享权重；或者深层网络中，每隔几层复用同一个线性变换。我做过一个跨语言翻译模型，源语言和目标语言的词表共享了80%的子词，这时候权重共享能大幅减少参数量。

来看一个跨层共享的经典模式：

class CrossLayerSharedNet(nn.Module):
    def __init__(self, vocab_size, embed_dim, num_layers):
        super().__init__()
        # 只创建一个Embedding层
        self.shared_embed = nn.Embedding(vocab_size, embed_dim)
        
        # 多个层都引用同一个embed
        self.layers = nn.ModuleList([
            nn.Sequential(
                self.shared_embed,  # 直接传对象引用
                nn.Linear(embed_dim, embed_dim),
                nn.ReLU()
            )
            for _ in range(num_layers)
        ])
        
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

这里有个容易踩的坑：如果你用nn.ModuleList包装多个层，每个层里都nn.Embedding(...)，那它们就是独立的。正确的做法是——把同一个Embedding对象传给多个层。

还有一种更隐蔽的跨层共享：转置权重共享。比如在自编码器中，编码器的最后一层权重和解码器的第一层权重互为转置。我做过一个文本降噪模型，用这种方式让参数量直接减半。

class TiedAutoEncoder(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super().__init__()
        self.encoder_weight = nn.Parameter(torch.randn(input_dim, hidden_dim))
        self.encoder_bias = nn.Parameter(torch.zeros(hidden_dim))
        self.decoder_bias = nn.Parameter(torch.zeros(input_dim))
        
    def forward(self, x):
        # 编码
        hidden = torch.relu(x @ self.encoder_weight + self.encoder_bias)
        # 解码时用转置权重
        recon = hidden @ self.encoder_weight.T + self.decoder_bias
        return recon

4.4 三种方式的对比

实现方式	代码量	可控性	推荐场景
手动实现	多	最高	需要精细控制梯度流、自定义查表逻辑
nn.Parameter共享	少	中	标准Embedding共享、同层多次调用
跨层共享	中	高	Encoder-Decoder共享、转置权重、深层网络复用

我个人建议：能用nn.Embedding对象引用解决的，就别手动搞。除非你有特殊需求，比如要对梯度做mask或者自定义查表索引。嗯，实际项目中90%的场景用第二种方式就够了。

最后提醒一句：权重共享虽然省参数，但也会限制模型的表达能力。我见过有人把整个Transformer所有层的FFN都共享了，结果模型怎么都训不上去。共享要适度，别为了省参数而牺牲性能。