算力直降97%，GPT-3存储只用20MB？！这篇直接在1.58-bit下训练模型的新论文火了
无需反向传播
一水 发自 凹非寺
量子位 | 公众号 QbitAI
好家伙！1750亿参数的GPT-3只需20MB存储空间了？！
基于1.58-bit训练，在不损失精度的情况下，大幅节省算力（↓97%）和存储（↓90%）。
最近，从事机器学习的Will小哥发了一篇论文，直接引来几十万网友or同行围观。
他提出了一项名为“noise_step”的新技术，允许模型直接在1.58-bit低精度下训练，且无需反向传播或动量（Momentum）加速，从而降低算力和存储消耗。
对此，网友们也纷纷发来祝贺，表示很高兴看到模型越来越具有性价比。
话不多说，来看论文具体内容。
反向传播不需要了
简单说，noise_step的目标是通过降低模型训练的精度要求，来减少算力和存储消耗。
一开始，作者提到前人研究已经表明，大语言模型（LLM）的推理可以在1.58-bit精度下进行，且不会有任何性能损失。
比如下面这篇论文，有人引入了1-bit的LLM变体（即BitNet b1.58），其中LLM的每个参数或权重都是三元的{-1, 0, 1}。
它在困惑度（perplexity）和最终任务性能上与全精度（FP16或BF16）的Transformer LLM相匹配，同时在延迟、存储、吞吐量和算力消耗方面成本更低。
然而，上述变体是在推理时使用低精度，而在训练时仍需高精度权重。
因此，noise_step的一个核心区别是：无需反向传播。
允许模型直接在1.58-bit（三元）精度下进行训练，而不需要传统的反向传播（从后向前检查每一层）或动量方法。
注：反向传播(Backpropagation)是训练神经网络的核心算法，它通过反向逐层计算损失函数对每个权重的梯度，来反向逐层更新网络的权重，从而最小化损失函数。
具体而言，will小哥参考了《Gradients without Backpropagation》这篇论文，其中介绍了雅可比向量积（Jacobian Vector Product，JVP）这种不依赖反向传播的梯度估计方法。
简单说，通过在前向传播中引入随机性，可以生成一个随机向量。这个随机向量与目标函数的梯度之间的对齐可以通过计算JVP来评估。
通过在多个随机方向上重复JVP计算，可以收集足够的信息来估计整个梯度向量，从而实现不依赖于反向传播的梯度估计。
will小哥的具体方法如下：
训练成本打下来了
上述方式意味着，noise_step允许使用更稀疏的随机向量和简单的对齐值。
要知道传统的梯度计算需要大量计算资源，而noise_step由于不需要存储或传输大量数据，从而减少了存储使用。
此外，由于noise_step使用伪随机噪声，它只需要一个种子（初始值）就能复现整个训练过程，这意味着不需要存储大量的扰动向量，从而进一步减少了存储需求。
而且使用noise_step训练的模型可以存储训练步骤而非权重，这可能会大幅缩小模型尺寸，从而更快地下载模型。
按照will小哥的说法，也许今后一秒钟内下载一个SOTA模型？
同时由于上面提到的伪随机噪声方法，这种特性允许恢复权重的完整历史，因为每个步骤都是确定性的，并且可以独立于其他步骤进行计算。
因此微调将变得更加高效，甚至可能允许对过去的训练步骤进行编辑，例如翻转（negation）或屏蔽（masking）。
举个例子，如果发现某个训练步骤对模型性能产生了负面影响，可以对其进行调整而不必重新训练整个模型。
也就是说，人们在训练过程中能进行更精细的控制和调整了。
最后，作者认为这种方式尤为适合分布式训练。
在分布式训练中，通常需要在不同的计算节点之间同步梯度和优化器状态，这会限制训练的速度。而noise_step通过减少每个扰动所需的位数，显著降低了通信量，从而提高了分布式训练的效率。
不过这也导致模型泄露变得更加容易，因为整个模型可以通过几个字节的训练步骤来传输。
对了，will小哥表示JVP可以和正常推理并行运行，几乎不增加成本。
除了论文，他也提供了一个CPU实现过程：
One More Thing
BTW，小哥在上分享完论文后，还顺带吐槽arXiv不给理由就拒绝了这篇论文。
后来他才补充，原来是卡在了背书（endorse）这一项，也就是需要现有用户的推荐或通过其他方式获得背书。
无奈之下，这篇论文目前被小哥放在了GitHub上。
感兴趣的童鞋可以进一步查看。
论文：
https://github.com/wbrickner/noise_step?tab=readme-ov-file
CPU实现过程：
https://colab.research.google.com/drive/1hXzf5xB4INzMUNTlAB8CI1V10-JV7zyg?usp=sharing
参考链接：
https://x.com/_brickner/status/1871348156786704657
- 智谱版o1终于也来了：直接拿下考研数学，一句话就能做小游戏！2024-12-31
- 你的下一个浏览器，可以是豆包电脑版！2024-12-25
- DeepSeek V3“报错家门”：我是ChatGPT2024-12-29
- 英特尔新款游戏显卡卖爆！24G大显存版也在路上，这下AI玩家也要抢？2024-12-24
相关阅读
