arXiv:2410.02117cs.LGstat.ML2024-10NeurIPS被引 9

提出连续空间搜索高效线性层,显著提升计算效率。

Searching for Efficient Linear Layers over a Continuous Space of Structured Matrices

  • 在连续结构空间中统一搜索线性算子,涵盖多种低秩结构
  • 发现参数共享少、秩高的结构具更好扩展性,提升计算效率
  • 设计BTT-MoE,每层独立稀疏化,比标准MoE更高效

密集线性层是大型神经网络的主要计算瓶颈,亟需更高效的替代方案。以往研究仅关注少量手工设计的结构,未考察这些结构在模型规模与训练样本最优分配下是否优于密集层的计算最优扩展规律。本文提出一个统一框架,可在所有可通过爱因斯坦求和表示的线性算子中进行搜索,涵盖低秩、克罗内克、张量-列车(Tensor-Train)、块张量-列车(BTT)、单峰(Monarch)等多种已有结构及众多新结构。我们基于计算与代数属性建立分类体系,发现计算最优扩展规律主要由两个变量决定:较小的ω(衡量参数共享)和较大的ψ(衡量秩)能带来更优表现。基于‘高参数密度结构性能最佳’的洞察,提出BTT-MoE——一种通过在BTT结构中稀疏化计算而得到的新颖专家混合(MoE)架构。与传统在整个前馈网络中稀疏化的MoE不同,BTT-MoE在模型每一层(包括注意力模块中的投影矩阵)均学习一个独立的MoE。实验表明,BTT-MoE相比密集层和标准MoE均实现了显著的计算效率提升。

原文摘要 · Abstract (English)

Dense linear layers are the dominant computational bottleneck in large neural networks, presenting a critical need for more efficient alternatives. Previous efforts focused on a small number of hand-crafted structured matrices and neglected to investigate whether these structures can surpass dense layers in terms of compute-optimal scaling laws when both the model size and training examples are optimally allocated. In this work, we present a unifying framework that enables searching among all linear operators expressible via an Einstein summation. This framework encompasses many previously proposed structures, such as low-rank, Kronecker, Tensor-Train, Block Tensor-Train (BTT), and Monarch, along with many novel structures. To analyze the framework, we develop a taxonomy of all such operators based on their computational and algebraic properties and show that differences in the compute-optimal scaling laws are mostly governed by a small number of variables that we introduce. Namely, a small $ω$ (which measures parameter sharing) and large $ψ$ (which measures the rank) reliably led to better scaling laws. Guided by the insight that full-rank structures that maximize parameters per unit of compute perform the best, we propose BTT-MoE, a novel Mixture-of-Experts (MoE) architecture obtained by sparsifying computation in the BTT structure. In contrast to the standard sparse MoE for each entire feed-forward network, BTT-MoE learns an MoE in every single linear layer of the model, including the projection matrices in the attention blocks. We find BTT-MoE provides a substantial compute-efficiency gain over dense layers and standard MoE.

线性层结构化矩阵MoE高效计算

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。