Flash Multi-Head Feed-Forward Network

2025-12-07 · Minshen Zhang, Xiang Hu, Jianguo Li, Wei Wu, Kewei Tu arxiv

We explore Multi-Head FFN (MH-FFN) as a replacement of FFN in the Transformer architecture, motivated by the structural similarity between single-head attention and FFN. While multi-head mechanisms enhance expressivity in attention, naively applying them to FFNs faces two challenges: memory consumption scaling with the head count, and an imbalanced ratio between the growing intermediate size and the fixed head dimension as models scale, which degrades scalability and expressive power. To address these challenges, we propose Flash Multi-Head FFN (FlashMHF), with two key innovations: an I/O-aware fused kernel computing outputs online in SRAM akin to FlashAttention, and a design using dynamically weighted parallel sub-networks to maintain a balanced ratio between intermediate and head dimensions. Validated on models from 128M to 1.3B parameters, FlashMHF consistently improves perplexity and downstream task accuracy over SwiGLU FFNs, while reducing peak memory usage by 3-5x and accelerating inference by up to 1.08x. Our work establishes the multi-head design as a superior architectural principle for FFNs, presenting FlashMHF as a powerful, efficient, and scalable alternative to FFNs in Transformers.

📄 PDF Abstract BibTeX arXiv:2512.06989

Code (0)

등록된 구현이 없습니다.

LangFlash: Feed-forward 3D Language Gaussian Splatting from Sparse Unposed Images

2026-05-22 · Yilong Liu, Wanhua Li, Chen Zhu-Tian, Hanspeter Pfister arxiv

We present LangFlash, a feed-forward framework for 3D Language Gaussian Splatting that reconstructs 3D scenes parameterized by Gaussian primitives enriched with language-aligned semantic features from sparse unposed mult…

Novel View Synthesis Scene Understanding 3D Reconstruction

Flash-Mono: Feed-Forward Accelerated Gaussian Splatting Monocular SLAM

2026-04-03 · Zicheng Zhang, Ke Wu, Xiangting Meng, Keyu Liu 외 arxiv

Monocular 3D Gaussian Splatting SLAM suffers from critical limitations in time efficiency, geometric accuracy, and multi-view consistency. These issues stem from the time-consuming $\textit{Train-from-Scratch}$ optimizat…

Flash3D: Feed-Forward Generalisable 3D Scene Reconstruction from a Single Image

2024-06-06 · Stanislaw Szymanowicz, Eldar Insafutdinov, Chuanxia Zheng, Dylan Campbell 외

In this paper, we propose Flash3D, a method for scene reconstruction and novel view synthesis from a single image which is both very generalisable and efficient. For generalisability, we start from a "foundation" model f…

3D Scene Reconstruction Depth Estimation GPU Monocular Depth Estimation+1

FlashSVD: Memory-Efficient Inference with Streaming for Low-Rank Models

2025-08-02 · Zishan Shao, Yixiao Wang, Qinsi Wang, Ting Jiang 외 arxiv

Singular Value Decomposition (SVD) has recently seen a surge of interest as a simple yet powerful tool for large language models (LLMs) compression, with a growing number of works demonstrating 20-80% parameter reduction…

Reducing the Cost of Dropout in Flash-Attention by Hiding RNG with GEMM

2024-10-10 · Haiyue Ma, Jian Liu, Ronny Krashinsky

Dropout, a network operator, when enabled is likely to dramatically impact the performance of Flash-Attention, which in turn increases the end-to-end training time of Large-Language-Models (LLMs). The main contributor to…