FFN融合:重新思考大语言模型中的顺序计算
机器学习
2025-03-25 v1
摘要
我们提出了FFN融合(FFN Fusion),这是一种用于减少大型语言模型中顺序计算的架构优化技术,通过识别和利用并行计算的自然机会。我们的关键洞察是, Feed-Forward网络(FFN)层的序列,特别是删除特定注意力层后剩余的FFN层,往往可以以最小的精度影响实现并行化。我们发展了一种原则性的方法,用于识别和融合此类序列,将其转换为显著减少推理延迟且保持模型行为的并行操作。将这些技术应用于Llama-3.1-405B-Instruct,我们创建了Llama-Nemotron-Ultra-253B-Base(Ultra-253B-Base),一个高效且即将公开的模型,实现了推理延迟加速1.71倍,每token成本降低35倍,同时在各类基准测试中保持强大的性能。通过对49B至253B参数的模型进行大量实验,我们证明FFN融合在大规模模型上效果更为显著,并且可以补充量化和剪枝等现有优化技术。最具新奇性的是,我们发现即使包含注意力层和FFN层的完整Transformer块有时也能实现并行化,这为神经网络架构设计开辟了新的方向。
引用
@article{arxiv.2503.18908,
title = {FFN Fusion: Rethinking Sequential Computation in Large Language Models},
author = {Akhiad Bercovich and Mohammad Dabbah and Omri Puny and Ido Galil and Amnon Geifman and Yonatan Geifman and Izhak Golan and Ehud Karpas and Itay Levy and Zach Moshe and Najeeb Nabwani and Tomer Ronen and Itamar Schen and Elad Segal and Ido Shahaf and Oren Tropp and Ran Zilberstein and Ran El-Yaniv},
journal= {arXiv preprint arXiv:2503.18908},
year = {2025}
}