在并行混合神经网络中平衡计算负载与表示表达力
计算与语言
2025-05-29 v2
摘要
注意力机制与状态空间模型在混合网络中以串行或并行方式结合时，能提供互补的优势。在混合串行流水线中，它们交替地将 Transformer 应用于输入，然后将其输出馈入 SSM。这导致各个组件出现空闲期，增加了端到端延迟并降低了吞吐量上限。在并行混合架构中，Transformer 与 SSM 独立并行运行，这些配对被级联，一个配对的输出构成下一个配对的输入。存在两个问题：(i) 利用这些独立分支本质上发散的输出来创建具有表达力的知识表示，以及 (ii) 在保持表示保真度的同时，平衡这些并行分支之间的计算负载。在本工作中，我们提出了 FlowHN，一种新颖的并行混合网络架构，它支持多种负载均衡策略，这是通过在两个分支之间适当分配输入 token 来实现的。FlowHN 的两个创新区分因素包括：一是感知 FLOP 的注意力与 SSM 分支间的动态 token 划分，在计算负载上实现了高效均衡；二是一种融合来自各个分支高度发散输出的方法，用于增强表示表达力。两者结合实现了更佳的 token 处理速度，避免了瓶颈，同时与其他竞争工作相比，准确率显著提升。我们在具有 135M、350M 和 1B 参数的模型上对自回归语言建模进行了全面实验。FlowHN 优于串行混合模型及其并行对应物，实现了高达 4 倍的每秒 token 数（TPS）和 2 倍更优的模型 FLOPs 利用率（MFU）。
引用
@article{arxiv.2505.19472,
title = {Balancing Computation Load and Representation Expressivity in Parallel Hybrid Neural Networks},
author = {Mohammad Mahdi Moradi and Walid Ahmed and Shuangyue Wen and Sudhir Mudur and Weiwei Zhang and Yang Liu},
journal= {arXiv preprint arXiv:2505.19472},
year = {2025}
}