arXiv:2505.19472cs.CL2025-05

提出FlowHN架构,让注意力与状态空间模型并行计算,兼顾速度与精度。

Balancing Computation Load and Representation Expressivity in Parallel Hybrid Neural Networks

  • 通过动态分配输入令牌实现计算负载均衡
  • 支持高达4倍的每秒处理词数,且模型算力利用率翻倍
  • 适合需要高速推理和高表达力的序列建模任务

将注意力机制与状态空间模型(SSM)以并行方式结合,可发挥各自优势。但并行结构面临两大挑战:一是如何融合两个分支差异较大的输出以构建丰富表示,二是如何在保持表示质量的前提下平衡计算负载。本文提出FlowHN架构,通过感知计算量的动态令牌分配,在注意力与SSM分支间实现高效负载均衡;同时设计方法融合差异显著的分支输出,提升表示表达能力。在135M、350M和1B参数的自回归语言建模任务上进行实验,结果表明,相比串行混合模型及其并行基线,FlowHN实现最高4倍的每秒处理词数(TPS),模型浮点运算利用率(MFU)提升2倍,显著提升效率与精度。

原文摘要 · Abstract (English)

Attention and State-Space Models (SSMs) when combined in a hybrid network in sequence or in parallel provide complementary strengths. In a hybrid sequential pipeline they alternate between applying a transformer to the input and then feeding its output into a SSM. This results in idle periods in the individual components increasing end-to-end latency and lowering throughput caps. In the parallel hybrid architecture, the transformer operates independently in parallel with the SSM, and these pairs are cascaded, with output from one pair forming the input to the next. Two issues are (i) creating an expressive knowledge representation with the inherently divergent outputs from these separate branches, and (ii) load balancing the computation between these parallel branches, while maintaining representation fidelity. In this work we present FlowHN, a novel parallel hybrid network architecture that accommodates various strategies for load balancing, achieved through appropriate distribution of input tokens between the two branches. Two innovative differentiating factors in FlowHN include a FLOP aware dynamic token split between the attention and SSM branches yielding efficient balance in compute load, and secondly, a method to fuse the highly divergent outputs from individual branches for enhancing representation expressivity. Together they enable much better token processing speeds, avoid bottlenecks, and at the same time yield significantly improved accuracy as compared to other competing works. We conduct comprehensive experiments on autoregressive language modeling for models with 135M, 350M, and 1B parameters. FlowHN outperforms sequential hybrid models and its parallel counterpart, achieving up to 4* higher Tokens per Second (TPS) and 2* better Model FLOPs Utilization (MFU).

混合模型并行计算算力优化语言建模

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