arXiv:2506.13541cs.CLcs.LG2025-06EMNLP被引 1

动态分配注意力资源,让每个词都高效参与计算。

Mixture of Weight-shared Heterogeneous Group Attention Experts for Dynamic Token-wise KV Optimization

  • 按词重要性动态分配计算资源,不丢弃任何词。
  • 在相同内存预算下,困惑度更低,生成质量更高。
  • 适合需要高效推理的大模型部署场景。

Transformer 模型在因果语言建模中因键值(KV)缓存内存分配低效而面临可扩展性挑战,加剧了计算与存储压力。现有方法如分组查询注意力(GQA)和词级 KV 优化虽提升效率,但依赖固定资源分配,常丢弃“低优先级”词或静态分组,无法应对词重要性的动态变化。本文提出 mixSGA,一种新型混合专家(MoE)方法,实现词级计算与内存分配的动态优化。不同于以往方法,mixSGA 保留所有词,通过学习的重要程度评分动态路由至具有不同 KV 分组大小的专用专家,平衡粒度与效率。核心创新包括:(1) 基于学习重要性分数的词级专家选择路由机制,实现无丢弃的比例化资源分配;(2) 通过共享分组注意力投影实现参数开销最小化;(3) 引入辅助损失确保训练-推理一致性的一对一路由决策。在 Llama3、TinyLlama、OPT 及 Gemma2 等模型族上的广泛评估显示,mixSGA 显著优于静态基线。在指令跟随与持续预训练任务中,相同 KV 预算下,mixSGA 实现更高的 ROUGE-L 得分与更低的困惑度。

原文摘要 · Abstract (English)

Transformer models face scalability challenges in causal language modeling (CLM) due to inefficient memory allocation for growing key-value (KV) caches, which strains compute and storage resources. Existing methods like Grouped Query Attention (GQA) and token-level KV optimization improve efficiency but rely on rigid resource allocation, often discarding "low-priority" tokens or statically grouping them, failing to address the dynamic spectrum of token importance. We propose mixSGA, a novel mixture-of-expert (MoE) approach that dynamically optimizes token-wise computation and memory allocation. Unlike prior approaches, mixSGA retains all tokens while adaptively routing them to specialized experts with varying KV group sizes, balancing granularity and efficiency. Our key novelties include: (1) a token-wise expert-choice routing mechanism guided by learned importance scores, enabling proportional resource allocation without token discard; (2) weight-sharing across grouped attention projections to minimize parameter overhead; and (3) an auxiliary loss to ensure one-hot routing decisions for training-inference consistency in CLMs. Extensive evaluations across Llama3, TinyLlama, OPT, and Gemma2 model families show mixSGA's superiority over static baselines. On instruction-following and continued pretraining tasks, mixSGA achieves higher ROUGE-L and lower perplexity under the same KV budgets.

TransformerKV缓存动态路由MoE

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