Mixture-of-Channels: 利用稀疏FFN实现高效LLM预训练与推理
机器学习
2025-11-13 v1
摘要
大型语言模型(LLM)在多样化的人工智能任务中展现出卓越成功, 这与模型规模与训练数据与性能提升的相关规模定律有关. 然而, 这种规模化范式带来了显著的内存开销, 在训练与推理中构成重大挑战. 现有研究主要解决参数和优化器状态的内存降低问题, 但当FlashAttention实现后, 激活内存(尤其是前馈网络FFN)已成为关键瓶颈. 本文对LLM进行详细的内存剖析, 识别FFN激活为激活内存开销的主要来源. 基于此, 我们引入Mixture-of-Channels(MoC), 一种新颖的FFN架构, 通过SwiGLU的本机门控机制按token选择激活最相关的Top-K通道. MoC在预训练期间显著降低激活内存, 并通过部分权重加载到GPU SRAM来提升推理效率. 大量实验验证了MoC在保持竞争模型性能的同时, 能实现显著的内存节省与吞吐量提升.
引用
@article{arxiv.2511.09323,
title = {Mixture-of-Channels: Exploiting Sparse FFNs for Efficient LLMs Pre-Training and Inference},
author = {Tong Wu and Yutong He and Bin Wang and Kun Yuan},
journal= {arXiv preprint arXiv:2511.09323},
year = {2025}
}