Balancing Computation Load and Representation Expressivity in Parallel Hybrid Neural Networks

注意力机制和状态空间模型（SSM）在混合网络中以序列或并行方式结合时，可以提供互补的优势。在混合顺序管道中，它们交替地将输入应用于变压器，然后将其输出馈送到 SSM 中。这会导致单个组件的空闲时间增加，从而增加端到端延迟并降低吞吐量上限。在并行混合架构中，变压器与 SSM 独立并行运行，并且这些对被级联，其中一对的输出形成下一组的输入。两个问题是：(i) 在这些分离分支的固有分歧输出之间创建一种富有表现力的知识表示，以及 (ii) 在这些并行分支之间进行计算负载均衡，同时保持表示的保真度。在这项工作中，我们提出了 FlowHN，这是一种新颖的并行混合网络架构，能够通过在两个分支之间适当地分配输入标记来实现多种负载均衡策略。FlowHN 的两个创新性区别因素包括：一种基于 FLOP 的动态标记分配方法，该方法在注意力分支和 SSM 分支之间实现了高效的计算负载平衡；其次，一种融合各分支高度分歧输出的方法，用于增强表示的表达能力。它们共同实现了更快的标记处理速度，避免了瓶颈，同时相较于其他竞争方法显著提高了准确性。我们在具有 1.35 亿、3.5 亿和 10 亿参数的自回归语言建模任务上进行了全面实验。FlowHN 在性能上超越了顺序混合模型及其并行对应模型，实现了高达 4 倍的每秒标记数（TPS）和 2 倍的模型浮点运算利用率（MFU）。

本论文试图解决如何高效结合注意力机制（Attention）和状态空间模型（SSMs）以提升序列建模任务中的吞吐量与准确性问题。具体来说，论文关注如何减少计算瓶颈、优化并行架构中的负载均衡以及增强输出表示的表达能力。这是一个在混合模型架构设计领域的重要研究方向，但尚未有完美的解决方案。

论文提出了一种名为FlowHN的新型并行混合网络架构。其关键思路包括：1) 动态分配输入令牌到注意力和SSM分支，基于FLOP感知策略实现计算负载的高效平衡；2) 设计一种融合方法来处理两个分支产生的高度分歧输出，从而增强表示的表达性。这种设计不仅避免了计算瓶颈，还显著提高了模型的吞吐量和准确性。

实验表明，FlowHN在自回归语言建模任务中表现出色，相较于串行和并行混合模型，分别实现了高达4倍的Tokens per Second (TPS) 和2倍的Model FLOPs Utilization (MFU) 提升。此外，论文针对不同规模的模型（参数量分别为135M、350M和1B）进行了全面实验验证。虽然论文未明确提及代码开源，但其提出的动态负载均衡和输出融合技术值得进一步研究和扩展。

相关研究包括：1) “Combining Attention and SSMs for Sequence Modeling”，探讨了串行和并行混合架构的设计；2) “Efficient Transformers: A Survey”，综述了Transformer模型的效率优化方法；3) “State-Space Models for Neural Sequence Processing”，深入分析了SSMs在序列任务中的应用。这些研究为FlowHN的设计提供了理论基础和技术启发。