FFN Fusie: Heroverweging van Sequentieel Rekenen in Grote Taalmodellen

FFN Fusion: Rethinking Sequential Computation in Large Language Models

March 24, 2025

Auteurs: Akhiad Bercovich, Mohammad Dabbah, Omri Puny, Ido Galil, Amnon Geifman, Yonatan Geifman, Izhak Golan, Ehud Karpas, Itay Levy, Zach Moshe, Najeeb Nabwani, Tomer Ronen, Itamar Schen, Elad Segal, Ido Shahaf, Oren Tropp, Ran Zilberstein, Ran El-Yaniv

cs.AI

Samenvatting

We introduceren FFN Fusion, een architectonische optimalisatietechniek die sequentiële berekeningen in grote taalmodelen vermindert door natuurlijke mogelijkheden voor parallelisatie te identificeren en te benutten. Onze belangrijkste inzicht is dat sequenties van Feed-Forward Network (FFN) lagen, met name die welke overblijven na het verwijderen van specifieke aandachtslagen, vaak kunnen worden geparalleliseerd met minimale impact op de nauwkeurigheid. We ontwikkelen een principiële methodologie voor het identificeren en samenvoegen van dergelijke sequenties, waarbij ze worden omgezet in parallelle bewerkingen die de inferentielatentie aanzienlijk verminderen terwijl het modelgedrag behouden blijft. Door deze technieken toe te passen op Llama-3.1-405B-Instruct, creëren we Llama-Nemotron-Ultra-253B-Base (Ultra-253B-Base), een efficiënt en binnenkort publiekelijk beschikbaar model dat een 1.71X versnelling in inferentielatentie en 35X lagere kosten per token bereikt, terwijl het sterke prestaties behoudt op benchmarks. Door uitgebreide experimenten op modellen van 49B tot 253B parameters, tonen we aan dat FFN Fusion steeds effectiever wordt op grotere schaal en kan complementair zijn aan bestaande optimalisatietechnieken zoals kwantisatie en pruning. Het meest intrigerend is dat we ontdekken dat zelfs volledige transformerblokken die zowel aandacht als FFN lagen bevatten soms kunnen worden geparalleliseerd, wat nieuwe richtingen suggereert voor het ontwerp van neurale architecturen.

English

We introduce FFN Fusion, an architectural optimization technique that reduces sequential computation in large language models by identifying and exploiting natural opportunities for parallelization. Our key insight is that sequences of Feed-Forward Network (FFN) layers, particularly those remaining after the removal of specific attention layers, can often be parallelized with minimal accuracy impact. We develop a principled methodology for identifying and fusing such sequences, transforming them into parallel operations that significantly reduce inference latency while preserving model behavior. Applying these techniques to Llama-3.1-405B-Instruct, we create Llama-Nemotron-Ultra-253B-Base (Ultra-253B-Base), an efficient and soon-to-be publicly available model that achieves a 1.71X speedup in inference latency and 35X lower per-token cost while maintaining strong performance across benchmarks. Through extensive experiments on models from 49B to 253B parameters, we demonstrate that FFN Fusion becomes increasingly effective at larger scales and can complement existing optimization techniques like quantization and pruning. Most intriguingly, we find that even full transformer blocks containing both attention and FFN layers can sometimes be parallelized, suggesting new directions for neural architecture design.