PyTorch实现ALBERT模型时如何在编码器间共享权重参数

你当前的写法是循环创建了num_encoder个独立的EncoderBlock实例，每个实例参数完全独立，没有共享效果。ALBERT的参数共享核心逻辑是复用同一个模块实例，不会为每个编码器层生成独立参数，根据共享范围不同，实现方式如下：

• 全参数共享（ALBERT论文默认配置）
  所有编码器层的注意力、FFN、LayerNorm参数完全共用，有两种稳妥实现方式：
  第一种是直接复用同一个EncoderBlock实例传入序列容器：

  # 仅实例化1个编码器块，所有层共用这一份参数
  shared_block = EncoderBlock(
      hidden_size,
      num_head,
      self.use_leakyrelu,
      dropout_p,
  )
  self.encoder = MySequential(
      *[shared_block for _ in range(num_encoder)]
  )
  

  只要你的MySequential是基于nn.Module实现的，重复传入同一个实例不会生成多份参数，PyTorch会自动识别同一个模块对象，仅注册一次参数。

  如果担心自定义MySequential有兼容问题，可以用更直观的循环前向写法，不需要依赖Sequential容器：

  # 初始化时仅存1个共享块和编码器层数
  self.shared_encoder_block = EncoderBlock(
      hidden_size,
      num_head,
      self.use_leakyrelu,
      dropout_p,
  )
  self.num_encoder = num_encoder
  

  对应前向传播逻辑里循环调用同一个块N次：

  def forward(self, x, attn_mask=None, **kwargs):
      for _ in range(self.num_encoder):
          x = self.shared_encoder_block(x, attn_mask=attn_mask, **kwargs)
      return x
  

• 部分参数共享
  如果要实现论文中提到的仅共享注意力层、仅共享FFN层这类策略，只需要把需要共享的子模块单独抽离成全局唯一实例，传入每个编码器块即可。以仅共享多头注意力层为例：

  # 先实例化要全局共享的注意力层
  shared_attn = MultiHeadAttention(hidden_size, num_head, dropout_p)
  self.encoder = MySequential(
      *[EncoderBlock(
          hidden_size,
          num_head,
          self.use_leakyrelu,
          dropout_p,
          shared_attn_module=shared_attn # 修改EncoderBlock构造逻辑，支持传入外部共享的注意力模块
        ) for _ in range(num_encoder)]
  )
  

  每个编码器块的FFN、LayerNorm保持独立初始化，仅注意力层复用同一份参数。

写完可以用两行代码验证共享是否生效，打印不同层同名参数的内存地址，地址完全一致就说明共享成功：

print(id(self.encoder[0].attn.qkv.weight))
print(id(self.encoder[1].attn.qkv.weight))

注意不要用copy.deepcopy复制编码器块，深拷贝会生成独立的参数副本，无法实现共享。

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