将注意力分段分配给专家,提升大模型计算效率
Optimal Scaling Laws for Efficiency Gains in a Theoretical Transformer-Augmented Sectional MoE Framework
- 按嵌入维度分段路由,而非整块令牌分配
- 推导出专家数量与模型规模的最优非线性关系
- 为高效扩展大模型提供理论指导,适合架构设计者
本文提出一种理论框架,构建基于Transformer增强的分段式混合专家(Sectional MoE)结构,旨在提升计算效率并保持模型可扩展性。与传统MoE不同,该方法不将完整令牌嵌入路由至专家,而是将嵌入维度拆分,将每个令牌表示的片段分配给专用专家。为缓解表示损失,引入预专家Transformer层,在令牌间重新计算注意力,并降低序列长度维度。通过扩展理论,推导出专家数量与模型维度、序列长度及系统开销之间的非线性最优缩放规律,获得在特定架构与硬件约束下可解析或数值求解的最优专家数表达式。该框架不仅提供不同架构下计算效率的理论边界,还指导大模型实际扩展设计。尽管尚无实证验证,本文提出了未来评估框架效率、可扩展性与实用性的全面实验路线。
原文摘要 · Abstract (English)
This paper introduces a theoretical framework for a Transformer-augmented, sectional Mixture-of-Experts (MoE) architecture that aims to enhance computational efficiency while preserving model scalability. Unlike conventional MoE models, which route entire token embeddings to selected experts, our approach portions the embedding dimension itself -- assigning segments of each token's representation to dedicated experts. To combat losses in token representation, we utilize a pre-expert transformer layer to recompute attention across tokens and reduce the sequence length dimensionality. We extend our theory by deriving optimal scaling laws that a non-linear relationship between the number of experts and factors such as model dimensionality, sequence length, and system overhead. These formulations yield closed-form and numerically-solvable expressions for identifying the optimal expert count under given architectural and hardware constraints. As a result, our framework not only provides theoretical bounds for computing efficiency with varying frameworks but also guides practical design choices for scaling large models effectively. While empirical validation is pending, we present a comprehensive experimental road map to evaluate the framework's efficiency, scalability, and practicality in future work.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。