用于语言模型内存高效适配的赢者通吃行列采样
机器学习
2024-11-08 v4 计算与语言
摘要
随着模型规模的快速增长，由于极大的内存占用，微调大型预训练语言模型已变得愈发困难。以往工作通常聚焦于减少网络中可训练参数的数量。尽管模型参数确实贡献了内存使用，但训练期间的主要内存瓶颈来自存储特征图（亦称激活值），因为它们对梯度计算至关重要。值得注意的是，神经网络通常使用随机梯度下降进行训练。我们认为，在随机优化中，只要梯度估计量无偏且方差合理，模型便能处理含噪梯度。遵循这一动机，我们提出了一类称为 WTA-CRS 的新无偏估计量，用于降低方差的矩阵乘积，其仅需存储子采样后的激活值以计算梯度。我们的工作从理论与实验上证明，在微调 transformers 的背景下，我们所提估计量相较于现有估计量表现出更低的方差。通过在 transformers 中以我们的近似线性运算替代原线性运算，我们可实现高达 2.7 的峰值内存削减且几乎无精度损失，并支持高达 更大的批大小。在相同硬件下，WTA-CRS 通过应用更大模型及/或借助更大批大小实现更快训练速度，从而获得更好的下游任务性能。
引用
@article{arxiv.2305.15265,
title = {Winner-Take-All Column Row Sampling for Memory Efficient Adaptation of Language Model},
author = {Zirui Liu and Guanchu Wang and Shaochen Zhong and Zhaozhuo Xu and Daochen Zha and Ruixiang Tang and Zhimeng Jiang and Kaixiong Zhou and Vipin Chaudhary and Shuai Xu and Xia Hu},
journal= {arXiv preprint arXiv:2305.15265},
year = {2024}
}