提出BASIS算法,用压缩张量实现高效反向传播,大幅降低内存占用。
BASIS: Balanced Activation Sketching with Invariant Scalars for "Ghost Backpropagation"

- 用低秩张量压缩梯度计算,解耦激活内存与批大小和序列长度
- 在R=32时验证损失仅略高于精确反向传播(6.575 vs 6.616)
- 支持极端压缩(R=1)仍稳定收敛,适合大模型训练优化
精确反向传播所需的激活内存随网络深度、上下文长度和特征维度线性增长,形成O(L * BN)的空间瓶颈(B为序列批大小,N为特征维度),长期制约深度神经网络的扩展。尽管随机自动微分试图缓解此问题,但常遭遇灾难性方差。本文提出BASIS(平衡激活草图与不变标量),一种高效的反向传播算法,完全解耦激活内存与批量和序列维度。BASIS传播精确误差信号(dX)以保证梯度流完整,但使用大规模压缩的秩-R张量计算权重更新(dW)。为解决草图梯度的根本不稳定性,提出两种新机制:平衡哈希严格消除非对角碰撞方差,不变标量通过合理偏差-方差权衡,确定性地保持空间几何的连续能量范数。理论上,BASIS将激活内存降至O(L * RN),显著减少反向传播中的矩阵乘法开销。实验上,对GPT架构训练50,000步验证理论保证:当R=32时,BASIS达到与精确反向传播相当的验证损失(6.575对比6.616),并略有超越,具有隐式正则化效果。值得注意的是,稳定的幅值轨迹使模型在极端空间压缩(R=1)下仍能平滑收敛,证明估计器的极高鲁棒性。代码已公开于https://github.com/VladimerKhasia/basis。
原文摘要 · Abstract (English)
The activation memory required for exact backpropagation scales linearly with network depth, context length, and feature dimensionality, forming an O(L * BN ) spatial bottleneck (where B is the sequence-batch cardinality and N is the feature dimension). This constraint historically throttles the scaling of deep neural networks. While randomized automatic differentiation attempts to mitigate this, it historically suffers from catastrophic variance. In this paper, we introduce BASIS (Balanced Activation Sketching with Invariant Scalars), an efficient backpropagation algorithm that fully decouples activation memory from the batch and sequence dimensions. BASIS propagates the exact error signal (dX) to preserve flawless gradient flow, but computes the weight updates (dW) using massively compressed rank-R tensors. To solve the foundational instability of sketched gradients, we propose two novel mechanisms: Balanced Hashing, which strictly eliminates off-diagonal collision variance, and Invariant Scalars, a principled bias-variance tradeoff that deterministically preserves the exact continuous energy norm of the spatial geometry. Theoretically, BASIS reduces activation memory to O(L * RN ) and heavily decreases the backward pass matrix-multiplication footprint. Empirically, training a GPT architecture for 50,000 steps validates our theoretical guarantees: at R = 32, BASIS achieves parity with (and marginally outperforms) exact backpropagation validation loss (6.575 vs. 6.616), acting as an implicit regularizer. Remarkably, the stabilized magnitude trajectory allows the model to converge smoothly even under extreme spatial compression (R = 1), proving the extreme robustness of the estimator. The code is available at https://github.com/VladimerKhasia/basis
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。