2023_NIPS_Winner-Take-All Column Row Sampling for Memory Efficient Adaptation of Language Model

文章总结与翻译

一、主要内容

该研究针对大型预训练语言模型微调时的内存瓶颈问题展开，核心聚焦于训练过程中占比超70%的激活值（特征图）存储开销。现有参数高效微调方法（如LoRA、Adapter）仅减少优化器状态内存，未解决激活值存储这一主要瓶颈。

研究提出基于随机优化的核心特性——模型可容忍带合理方差的无偏梯度，设计了一种低方差无偏估计器WTA-CRS（Winner-Take-All Column-Row Sampling）。其核心思路是通过列-行采样（CRS）对矩阵乘法进行近似，仅存储子采样后的激活值用于梯度计算，同时通过区分高概率区域（确定性求和）与低概率区域（随机采样）降低估计方差。

实验验证显示，WTA-CRS在Transformer模型中替换线性操作后，可实现最高2.7倍的峰值内存 reduction，且几乎无精度损失；支持最高6.4倍的批次大小提升，使T5-3B等大模型能在24GB显存GPU（如RTX3090Ti）上完成微调。该方法与现有参数高效微调技术正交，可组合使用进一步优化内存效率，同时提升训练吞吐量最高达1.21倍。

二、创新点

1. 低方差无偏估计器设计：提出WTA-CRS估计器，通过划分高/低概率列-行对区域，结合确定性求和与随机采样，在保持无偏性的同时降低方差，理论上证明其方差优于传统CRS方法。
2. 激活值内存优化：针对Transformer线性层的矩阵乘法，仅在反向传播阶段应用近似计算，避免前向传播非线性激活导致的偏差，从根源减少激活值存储开销。