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

要約

密な線形層は、大規模なニューラルネットワークにおいて支配的な計算ボトルネックであり、より効率的な代替物の重要な必要性を提示している。これまでの取り組みでは、少数の手作業で構造化された行列に焦点が当てられており、モデルサイズと学習例の両方が最適に配分されている場合に、これらの構造が計算最適スケーリング則の点で密な層を上回ることができるかどうかの調査がおろそかにされていた。本論文では、アインシュタイン和で表現可能な全ての線形演算子を検索できる統一的な枠組みを提示する。このフレームワークは、低ランク、クロネッカー、Tensor-Train、Block Tensor-Train(BTT)、Monarchなど、以前に提案された多くの構造と、多くの新しい構造を包含する。このフレームワークを解析するために、計算量と代数的性質に基づく全ての演算子の分類法を開発し、計算最適スケーリング則の違いは、我々が導入した少数の変数によってほとんど支配されることを示す。すなわち、小さい$omega$（パラメータ共有を測定する）と大きい$ppsi$（ランクを測定する）は、より良いスケーリング則を確実に導いた。単位計算量あたりのパラメータを最大化するフルランク構造が最適であるという洞察に導かれ、BTT-MoEを提案する。BTT-MoEは、BTT構造において計算をスパース化することによって得られる新しいMoE（Mixture-of-Experts）アーキテクチャである。各フィードフォワードネットワーク全体に対する標準的なスパースMoEとは対照的に、BTT-MoEは、注目ブロックにおける射影行列を含む、モデルの全ての線形層におけるMoEを学習する。BTT-MoEは、密な層や標準的なMoEに比べて、計算効率が大幅に向上することがわかった。

要約(オリジナル)

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 $\omega$ (which measures parameter sharing) and large $\psi$ (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.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL