权重共享在模型压缩中的应用ALBERT的跨层参数共享与通用化分析一、参数共享的直觉与ALBERT的激进设计语言模型的参数规模在过去五年增长了四个数量级从BERT-base的110M到GPT-4的约1.7T但其中相当比例的参数是冗余的。一项被广泛引用的分析表明BERT的12层Transformer中相邻层之间的注意力权重模式具有高度相似性——第5层的注意力头和第6层的对应头经常学习到相似的语法模式。ALBERTA Lite BERT, Lan et al., 2020基于这一观察提出了一个激进的设计跨层参数共享。ALBERT的所有Transformer层共享同一组参数——这意味着12层ALBERT-base和24层ALBERT-large的参数量完全相同因为只有一组参数被复制了12或24次。与传统BERT相比ALBERT-base的参数量从110M降至12M约9倍压缩同时在下游任务上保持了近90%的性能。这一设计的核心问题是如果12层都做相同的事那为什么需要12层ALBERT的回答是虽然参数共享了但每一层的输入嵌入不同——在第1层处理的表示与第12层处理的表示完全不同。参数共享迫使模型学习通用的特征变换这些变换在不同抽象层次上都有用。二、参数共享的实现方式与粒度选择跨层参数共享的实现粒度可以分为三个层次全层共享ALBERT的方案所有Transformer层的所有参数自注意力投影矩阵、FFN权重、LayerNorm参数全部共享。这是最激进的方案压缩率最高9×但性能损失也最大约3-5%的GLUE基准下降。仅注意力共享共享自注意力的Q/K/V/O投影矩阵但FFN层保持独立。这基于一个假设不同层的注意力模式相似句法级的注意力在学习早期就已形成但FFN的特征变换是层次化的从低级特征到高级语义。仅FFN共享与上一种方案对称——让每层有独立的注意力模式但共享特征变换函数。这种方案的动机是不同层需要关注不同的token间关系但将关注结果转化为有用特征的方式可以是通用的。 三种参数共享粒度的PyTorch实现对比 import torch import torch.nn as nn import copy class SharedTransformerLayer(nn.Module): 单层Transformer可被多次复用以实现跨层参数共享 def __init__(self, hidden_size: int 768, num_heads: int 12, ffn_size: int 3072): super().__init__() self.attention nn.MultiheadAttention( hidden_size, num_heads, batch_firstTrue ) self.ffn nn.Sequential( nn.Linear(hidden_size, ffn_size), nn.GELU(), nn.Linear(ffn_size, hidden_size), ) self.ln1 nn.LayerNorm(hidden_size) self.ln2 nn.LayerNorm(hidden_size) def forward(self, x: torch.Tensor) - torch.Tensor: # Pre-Norm 自注意力 残差 attn_out, _ self.attention(self.ln1(x), self.ln1(x), self.ln1(x)) x x attn_out # Pre-Norm FFN 残差 x x self.ffn(self.ln2(x)) return x class ALBERTLike(nn.Module): 类ALBERT模型支持多种参数共享粒度 def __init__( self, vocab_size: int 30000, hidden_size: int 768, num_layers: int 12, num_heads: int 12, ffn_size: int 3072, sharing_mode: str all, # all | attention_only | ffn_only | none ): Args: sharing_mode: 参数共享模式 - all: 所有层共享全部参数ALBERT方案 - attention_only: 仅共享注意力层FFN独立 - ffn_only: 仅共享FFN层注意力独立 - none: 每层独立参数传统BERT super().__init__() self.embedding nn.Embedding(vocab_size, hidden_size) self.num_layers num_layers self.sharing_mode sharing_mode if sharing_mode all: # 1组共享参数复用于所有层 self.shared_layer SharedTransformerLayer( hidden_size, num_heads, ffn_size ) self.layers [self.shared_layer] * num_layers elif sharing_mode attention_only: # 所有层共享1组注意力参数 self.shared_attention nn.MultiheadAttention( hidden_size, num_heads, batch_firstTrue ) self.shared_ln1 nn.LayerNorm(hidden_size) # 每层独立的FFN self.ffns nn.ModuleList([ nn.Sequential( nn.Linear(hidden_size, ffn_size), nn.GELU(), nn.Linear(ffn_size, hidden_size), ) for _ in range(num_layers) ]) self.ln2s nn.ModuleList([ nn.LayerNorm(hidden_size) for _ in range(num_layers) ]) elif sharing_mode ffn_only: # 所有层共享1组FFN参数 self.shared_ffn nn.Sequential( nn.Linear(hidden_size, ffn_size), nn.GELU(), nn.Linear(ffn_size, hidden_size), ) self.shared_ln2 nn.LayerNorm(hidden_size) # 每层独立的注意力 self.attentions nn.ModuleList([ nn.MultiheadAttention(hidden_size, num_heads, batch_firstTrue) for _ in range(num_layers) ]) self.ln1s nn.ModuleList([ nn.LayerNorm(hidden_size) for _ in range(num_layers) ]) else: # none self.layers nn.ModuleList([ SharedTransformerLayer(hidden_size, num_heads, ffn_size) for _ in range(num_layers) ]) def forward(self, input_ids: torch.Tensor) - torch.Tensor: x self.embedding(input_ids) if self.sharing_mode all: for _ in range(self.num_layers): x self.shared_layer(x) elif self.sharing_mode attention_only: for i in range(self.num_layers): ln_x self.shared_ln1(x) attn_out, _ self.shared_attention(ln_x, ln_x, ln_x) x x attn_out x x self.ffns[i](self.ln2s[i](x)) elif self.sharing_mode ffn_only: for i in range(self.num_layers): ln_x self.ln1s[i](x) attn_out, _ self.attentions[i](ln_x, ln_x, ln_x) x x attn_out x x self.shared_ffn(self.shared_ln2(x)) else: # none for layer in self.layers: x layer(x) return x def count_parameters(self) - dict: 统计不同共享模式下的参数量 total sum(p.numel() for p in self.parameters()) trainable sum(p.numel() for p in self.parameters() if p.requires_grad) return {total: total, trainable: trainable}三、参数共享的普适性何时有效、何时失效ALBERT的成功引出了一个自然的问题参数共享能否推广到其他模型和任务实验证据指向一个条件性的结论。有效场景在编码器架构的NLP理解任务GLUE、SQuAD中全层参数共享的ALBERT保持了BERT约90%的性能同时参数减少9倍。这表明在理解型任务中深度主要提供迭代精炼而非功能分化——相同的变换反复应用每次在上一层的输出上操作逐步精化表示。失效场景在生成任务如翻译、摘要中全层参数共享的性能损失显著更大约10-15%。这可能是因为生成任务需要更丰富的层次化特征——浅层处理句法、中层处理语义、深层处理语用——而这些不同层次的特征变换不能由同一组参数有效地完成。部分共享的折中将12层划分为3组浅层组/中层组/深层组组内共享参数、组间独立。这种分组共享策略在参数效率约3倍压缩和性能保持2%下降之间取得了更好的平衡在迁移学习场景中表现尤为稳健。四、参数共享与知识蒸馏的协同效应参数共享和知识蒸馏是两种正交的压缩策略——前者压缩模型参数的存储后者压缩模型的计算图通过学生-教师框架。两者的组合可以产生超线性大于各自独立效果的简单相加的压缩效果。TinyBERT的实验表明使用BERT-base作为教师、ALBERT架构作为学生进行蒸馏得到的模型在参数量约5MB和推理速度上均达到了可用级别同时保持了BERT-base约95%的GLUE性能。这种协同效应的一个可能解释是参数共享减少了学生模型的搜索空间所有层都执行相同的变换使得蒸馏过程中的优化更容易收敛——教师不需要指导12个不同的层各学习什么只需要指导一个通用的变换如何在不同深度上发挥作用。五、总结ALBERT的跨层参数共享通过将12组独立参数压缩为1组实现了9倍的参数量缩减。全层共享在理解任务中保持约90%的BERT性能但在生成任务中性能损失扩大仅注意力共享或仅FFN共享提供了参数量与性能之间的可调节折中分组共享层分组、组内共享是生产实践中更稳健的选择。参数共享作为一种独立于量化、剪枝、蒸馏的压缩维度其价值在于与其他压缩方法的正交性——共享蒸馏的组合可以产生比各自独立使用更大的压缩效果。在部署到资源受限环境移动端、浏览器WebAssembly时参数共享是少数几种不依赖特定硬件加速如INT8 Tensor Core即可显著缩减模型体积的策略之一。