用分阶段成对混合取代密集矩阵,实现近线性训练
Rethinking Dense Linear Transformations: Stagewise Pairwise Mixing (SPM) for Near-Linear Training in Neural Networks
- 将密集层替换为分阶段成对混合结构,降低计算复杂度
- 在O(nL)时间与参数下完成全局线性变换,支持精确前向后向传播
- 适合需要高效训练且结构化任务的模型,如注意力机制与循环网络
密集线性层是现代机器学习模型中主要的计算与参数开销来源,尽管其具有二次复杂度且常与学习表示的组合结构不匹配。本文提出分阶段成对混合器(SPM),一种结构化线性算子,通过稀疏成对混合阶段替代密集矩阵。一个SPM层以O(nL)的时间复杂度和参数量实现全局线性变换,其中L通常为常数或log₂n,且支持精确闭式前向与反向计算。SPM可作为前馈网络、循环架构、注意力机制等中的密集层即插即用替代品。我们推导了两种参数化形式的完整前向与反向表达:一种保持正交性的旋转型,另一种为通用2×2混合型。除了显著降低计算开销外,SPM的分阶段结构引入显式的组合归纳偏置,当与任务结构对齐时能限制模型容量并提升泛化性能。初步实验表明,在结构化学习任务上大幅降低实际运行时间并提升准确率,同时在真实基准上保持竞争力。
原文摘要 · Abstract (English)
Dense linear layers are a dominant source of computational and parametric cost in modern machine learning models, despite their quadratic complexity and often being misaligned with the compositional structure of learned representations. We introduce Stagewise Pairwise Mixers (SPM), a structured linear operator that replaces dense matrices with a composition of sparse pairwise-mixing stages. An SPM layer implements a global linear transformation in $O(nL)$ time with $O(nL)$ parameters, where $L$ is typically constant or $log_2n$, and admits exact closed-form forward and backward computations. SPM is designed as a drop-in replacement for dense linear layers in feedforward networks, recurrent architectures, attention mechanisms, etc. We derive complete forward and backward expressions for two parameterizations: an orthogonal norm-preserving rotation-based variant and a fully general $2 \times 2$ mixing variant. Beyond computational savings, the stagewise structure of SPM induces an explicit compositional inductive bias that constrains model capacity and improves generalization when aligned with task structure. We present proof-of-concept experiments demonstrating substantial reductions in wall-clock cost and improved accuracy on structured learning problems, while retaining competitive performance on real-world benchmarks.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。