用Top-K函数实现稀疏化
Q-Sparse最核心的操作是对输入张量应用Top-K稀疏化函数。
具体来说,Transformer架构在注意力层和前馈层中都使用nn.Linear线性层(矩阵乘法)进行投影,可以表示为Y=X·W^T。(其中X是输入张量,W代表其权重,Y为输出张量)
Q-Sparse中,对于输入激活张量X,首先计算其绝对值|X|并排序,找出绝对值最大的K个元素。
K是预先设定的超参数,决定稀疏化程度。
然后创建与X形状相同的二进制掩码张量M,对|X|中绝对值最大的K个元素对应位置,将M中相应位置设为1,其余位置设为0。
接着,将X与M进行Hadamard积(逐元素相乘)运算,得到稀疏化张量X_sparse。
前向传播时,X_sparse代替原始X参与后续计算(如矩阵乘法)。
由于X_sparse大部分元素为零,可显著减少计算量和内存带宽需求。
反向传播时,Q-Sparse使用直通估计器(STE)计算Top-K函数梯度。
STE通过直接将梯度传给稀疏化前的张量,避免梯度消失问题。
对于前馈层,Q-Sparse用平方ReLU函数代替常规ReLU,进一步提高激活稀疏性。
为适配量化模型,Q-Sparse在Top-K稀疏化前先对输入张量量化,确保与量化表示兼容。
对1-bit量化权重,Q-Sparse使用特定量化函数,α是权重张量W的平均绝对值。
60%激活参数达到相同效果
实验表明,Q-Sparse在稀疏率和模型表现上显著优于ReLU方法。
从头训练实验中,700M和7B的Llama模型使用70% top-K(40%整体稀疏率)可达到与密集baseline相当的训练损失。
继续训练实验中,Mistral-7B在激活参数为2.9B和3.8B时,各数据集得分未明显下降。
微调实验中,Qwen-7B和Mistral-7B用60%左右激活参数实现与密集模型接近的表现。
这意味着稀疏激活模型在推理时可显著减少激活参数,降低FLOPS消耗。
对量化模型,团队在BitNet b1.58上应用Q-Sparse并评估。结果显示Q-Sparse可无缝集成到量化模型中,不影响训练和收敛。
作者认为,Q-Sparse与量化技术结合可进一步提高大语言模型推理效率。
发现推理优化新"Scaling Law"
作者探究了模型性能、规模和稀疏率三者关系,有新发现。
稀疏激活模型的性能缩放定律:与密集模型类似,稀疏激活模型性能也遵循幂律缩放关系。
给定稀疏率S,模型收敛时损失函数值L(N,S)可近似为: