NVlabs/Sana模型中交叉注意力层偏置权值异常问题分析

背景概述

在深度学习的扩散模型领域，NVlabs推出的Sana系列模型(如Sana_1600M_1024px_MultiLing_diffusers)因其出色的多语言图像生成能力而备受关注。然而在模型微调过程中，研究者发现某些交叉注意力层(cross-attention)存在数值异常现象，特别是key投影层的偏置项(bias)出现了绝对值异常大的权重值(部分超过20)，这可能影响模型训练的稳定性。

问题现象深度解析

通过权重分析发现，在transformer模块的特定层(如第6、12、5、7等层的attn2.to_k.bias)存在显著的数值异常：

1. 数值范围异常：部分偏置项的最大绝对值达到20.06，最小值-19.98
2. 分布规律性：异常值主要集中在中间层(如block 3-14)，而极浅层和极深层相对正常
3. 结构特异性：问题主要出现在交叉注意力(attn2)的key投影层，而非query或value投影层

这种现象在模型微调时可能导致：

• 梯度爆炸风险：大权重值容易在反向传播时产生梯度不稳定
• 优化困难：损失函数曲面可能出现剧烈波动
• 特征表达扭曲：过大的偏置可能破坏注意力机制的正常运作

技术原理探究

交叉注意力层的key投影矩阵负责将条件输入(如文本嵌入)映射到key空间，其偏置项的异常增大可能源于：

1. 训练动态失衡：在长序列处理时，某些token获得了过大的注意力权重
2. 正则化不足：缺少适当的权重约束机制(如LayerNorm或QK Norm)
3. 初始化策略缺陷：偏置项初始化范围可能不适合深层架构

解决方案与改进

项目团队在后续的SANA-1.5版本中实施了针对性改进：

1. 引入QK标准化：对所有自注意力和交叉注意力层添加query-key归一化
2. 权重约束机制：通过参数初始化调整和正则化控制权重范围
3. 架构优化：重新平衡各层参数分布，特别关注中间层稳定性

实践建议

对于使用早期版本的研究者：

1. 微调策略：建议采用较小的学习率(如1e-6)和梯度裁剪
2. 监控重点：特别关注中间层(block 3-14)的梯度范数
3. 权重修复：可尝试对异常偏置进行渐进式缩放(scale down)

总结

大模型训练中的数值稳定性是保证微调效果的关键因素。NVlabs团队通过持续的架构改进，特别是引入注意力归一化机制，有效解决了交叉注意力层的权重异常问题。这为扩散模型的可控微调提供了重要参考，也体现了深度学习系统工程中数值稳定性的重要性。