# Balancing Computation Load and Representation Expressivity in Parallel Hybrid Neural Networks

arXiv: https://arxiv.org/abs/2505.19472 · HTML: https://mlsanity.com/paper/2505.19472/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 19 of its 21 references that are in the map.
- First posted on arXiv 26 May 2025.

## Builds on (most cited first, 19 of 19)

- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Phi-3 Technical Report: A Highly Capable Language Model Locally on Your Phone](https://mlsanity.com/paper/2404.14219/) (2024)
- [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://mlsanity.com/paper/2405.21060/) (2024)
- [RWKV: Reinventing RNNs for the Transformer Era](https://mlsanity.com/paper/2305.13048/) (2023)
- [Gated Linear Attention Transformers with Hardware-Efficient Training](https://mlsanity.com/paper/2312.06635/) (2023)
- [Hungry Hungry Hippos: Towards Language Modeling with State Space Models](https://mlsanity.com/paper/2212.14052/) (2022)
- [Jamba: A Hybrid Transformer-Mamba Language Model](https://mlsanity.com/paper/2403.19887/) (2024)
- [Repeat After Me: Transformers are Better than State Space Models at Copying](https://mlsanity.com/paper/2402.01032/) (2024)
- [Landmark Attention: Random-Access Infinite Context Length for Transformers](https://mlsanity.com/paper/2305.16300/) (2023)
- [Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling](https://mlsanity.com/paper/2406.07522/) (2024)
- [Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning Tasks](https://mlsanity.com/paper/2402.04248/) (2024)
- [Demystify Mamba in Vision: A Linear Attention Perspective](https://mlsanity.com/paper/2405.16605/) (2024)
- [Hymba: A Hybrid-head Architecture for Small Language Models](https://mlsanity.com/paper/2411.13676/) (2024)
- [Jamba-1.5: Hybrid Transformer-Mamba Models at Scale](https://mlsanity.com/paper/2408.12570/) (2024)
- [Bridging the Divide: Reconsidering Softmax and Linear Attention](https://mlsanity.com/paper/2412.06590/) (2024)
- [ZETA: Leveraging Z-order Curves for Efficient Top-k Attention](https://mlsanity.com/paper/2501.14577/) (2025)

## Abstract (arXiv)

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).