FFN Fusion: Rethinking Sequential Computation in Large Language Models
Authors: Akhiad Bercovich
Mohammed Dabbah
Omri Puny
Ido Galil
Amnon Geifman
Yonatan Geifman
Izhak Golan
Ehud Dov Karpas
Itay Levy
Zach Moshe
Najeeb Nabwani
Tomer Ronen
Itamar Schen
Ido Shahaf
Oren Tropp
Ran Zilberstein
Ran El-Yaniv
Venue: NeurIPS 2025 spotlight ()
Keywords: LLM
Architecture optimization
Inference efficiency
Primary area: deep_learning
TL;DR: An architectural optimization method that accelerate inference efficiency by replacing sequential computation in LLM with parallel computation while preserving accuracy.

Abstract: We introduce \textit{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 a 253B model (253B-Base), an efficient and soon-to-be publicly available model that achieves a 1.71$\times$ speedup in inference latency and 35$\times$ lower per-token cost while maintaining strong performance across benchmarks. 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.

FFN Fusion: Rethinking Sequential Computation in Large Language Models
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
Io Shahaf
Oren Tropp
Ran Zilberstein
Ran EI- Yaniy
NVIDIA, (abercovich,
mdabbah, opuny, relyaniv)Onvidia, com
Abstract
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 re- moval 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 (Ultra-
we create 253B-Base)) an efficient model that achieves a 1.71 x speedup in inference latency and 35x lower per token cost while maintaining strong performance across bench- marks. 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.
-lama-3.11005B-Insntte
Introduction
Large language models (LLMs) have emerged as of the most transformative technologies of our time, revolutionizing how we approach artificial intelligence and computation. From powering sophisticated virtual assistants [16] to enabling breakthrough scientific research 41 these models have evolved from academic curiosities
35 into indispensable tools that are reshaping entire industries. This transformation has been driven by rapid sealing to hundreds of billions parameters enabling unprecedented capabilities in reasoning, generation, and complex problem-solving 15 However, this extraordinary growth scale and capability comes at a critical cost: the computational demands of running these models have become a fundamental bottleneck. As these models push the boundaries of what possible with artificial intelligence, their deployment costs and resource requirements severely limit their accessibility, creating an urgent need for innovations that can make their capabilities widely available.
33
2
Research in LLM runtime optimization explored diverse approaches to managing computational demands. Traditional techniques like quantization 1 47 which reduce memory footprint and which
accelcrate inference through lower-precision arithmetic and pruning (2 [8
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

removes redundant parameters have become standard tools. A more recent innovation, Mixture- of of-Experts (MoE) 38 23 has demonstrated remarkable potential by dynamically activating only small subset of model parameters during inference. The DeepSeek- V3 architecture [6] pushes this approach to new extremes, employing 256 expert FFN modules at each layer while activating only 8 experts per token (plus one shared expert), effectively achieving the capabilities of much larger model while maintaining reasonable computational costs through sparse activation. However, each of these approaches faces distinct challenges: quantization encounters precision-accurac trade- offs at small bit numbers, pruning struggles to identify significant additional redundancy without compromising aceuracy performance (and pruning can be harnessed efficiency mainly when it is structured), and MoE architectures, while efficient for single-token inference or for very large batches, do not provide optimal throughput at small and intermediate batch sizes due to under- under-utilization of compute resources, as discussed in Appendix B Given the above limitations, there is an urgent need for complementary approaches that unlock new dimensions of efficiency improvement while maintaining the simplicity and predictable scaling characteristics of dense architectures.
In this work, we introduce two major contributions that fundamentally advance the state of LLM efficiency. First, we present FFN Fusion, a novel architectural optimization technique that challenges the conventional sequential nature of transformer computation, and is illustrated in Figure By iden- tifying and exploiting patterns of computational independence in FFN layers, our approach enables parallel execution across multiple GPUs while preserving model functionality. This parallclization is particularly effective on modern GPU nodes, where tensor-pest cansor-parallel implementations often suffer from synchronization delays between consecutive layers. By concentrating the computation into fewer layers and reducing cross-device communication, our method significantly improves hardware utilization. Our findings reveal that substantial portions of LLM computation can be parallelized with minimal accuracy impact, complementing existing runtime optimization techniques like pruning and quantization. Second, we demonstrate the practical impact these insights through Ultra-253B- Base, a state-of-the-art 253B parameter model derived from 1ama--.11400BB-- Llama-405B), derived using FFN Fusion and attention pruning. This model not only showcases the scalability of our approach, but also achieves remarkable efficiency gains while maintaining or exceeding its parent model's 's capabilities across a comprehensive suite of benchmarks.
(henceforth
FFN7
Attention
Step Apply Puzzle
FFN6
Attention
FFNS
FFN7
Step 2:
Attention
Attention
FFN Fusion
FFN4
FFN6
Attention
FFN5
[FFNZ, FFNB, FFN9]
FFN3
FFN4
Attention
Attention
FFN3
[FFN5, FFN6)
FFN2
FFN2
[FFN2 FFN3, FFN4]
Attention
Attention
Attention
FFN1
FFN1
FFN1
Attention
Attention
Attention
We apply Puzzle to partially remove FFN layers and remove entire attention layers. Step 2: We fuse consecutive FFN layers into single wide FFN layer. Fusion is denoted using the bracket notation |FFN ...., FFNi+k] ,FFN
Figure I: An overview of our FFN Fusion approach. Step 1:

---

The applicabilitt of FFN Fusion rests on fundamental insight about modern LLM architectures.
Recent work has shown that LLM display structural redundancy 44 larly attention mechanisms, which can be selectively removed with minimal accuracy impact (3
10
49
36
25
14,
particu-
often leaving models with extended sequences of consecutive FFN layers. We demonstrate that these FFN sequences exhibit surprisingly low inter-layer dependencies, enabling novel transformation: multiple sequential FFN layers can be fused into a single, wider layer that enables simple parallel execution. This transformation is particularly powerful, eliminating synchronization points while allowing for more efficient hardware utilization without requiring architectural redesign (as visualized in Figure 2) Through analysis and empirical validation, argue for the conditions under which this fusion preserves model behavior, showing that these conditions are commonly satisfied in practice, especially in larger models where the potential efficiency gains most significant. Most surprisingly, our preliminary investigations suggest that even complete transformer tion and FFN components- can sometimes be parallelized, pointing potential new directions in neural architecture design (see Appendix [B. Leveraging these insights, we develop Sitra-253B-Bar derived from Llama-405B through combination of attention pruning and FFN Fusion. This model achieves remarkable efficiency gains while maintaining or exceeding its parent model's ' capabilities:
containnn both atten-
A 1.71 x speedup in inference latency and 35x lower per-token cost at batch size 32.
Statc-of-the-an performance on key benchmarks, including 84.92% on Arena Hard, 86. 58% on HumanEval, 87.5 37.54% on MMLU Instruct, 72.25% on MMLUUPro, and 9.19 on MT- Bench.
Improving memory footprint with a 2x reduction in kv-cache memory, and reduced param- eter count from 405B to 253B.
Our key contributions are:
The discovery that some Feed-Forward Networks in transformer architectures can often be parallelized with minimal accuracy loss, challenging the conventional wisdom that sequential processing is necessary for these layers.
The creation of Ultra-253B-Ba Basse, a powerful 253B parameter model that matches or exceeds Llama-405B's 's capabilities while offering substantially improved efficiency, to be publicly released upon acceptance.
Preliminary evidence that even complete transformer flocks-contaiiing both attention and FFN components- can sometimes be parallelized, opening new possibilities for architectural innovations in large language models.
2 Preliminaries
The structure of L.LMs is typically based on the widely used trans- former architecture (4).. is generally composed of a series of sequential blocks, denoted as f!...., f each containing an attention layer and FFN layer. The attention layers capture the contextual relationships between input tokens, while the FFN layers apply transformations indepen- dently to each token. Given an input X E Raxd where n represent the sequence length and d the embedding dimension, a transformer block f Raxd + Rnxd is defined by the following equations:
Transformer-based LLMs.
g( x(X) X + (ttention(n) (X)) f(x) x + (x))))
where 7) is a token-level normalization module, defined by the equation T) (x.)) = Zt. + 8% with
I; E Rd being the i-th token, 8% E Rdd being a learnable scale factor, and 3 is the element-wise product. The FFN layer used in the models we experiment with is the SwiGLU module [37] which is defined as follows:
SwiGLU(X)=<rd = (xW o XWT)W3
where is the SiLU activation 20 function, W1. W2 E Rd xd. and W3 E xd Here, de represents the embedding dimension, and d, denotes the hidden dimension of the FFN. Other components of the LLM include an embedding function at the model's input, which maps tokens to vectorized embeddings, an additional normalization function, and linear function at the output, which normalizes the final block outputs and projects them to the vocabulary dimension.
3

---

Puzzle. Puzzle 33 is a neural architecture search (NAS) framework that optimizes trained LLM for inference efficiency. Starting with pre-trained model, it prunes or reconfigures each transformer block-ofter reducing or removing attention whhree preserving model quality through distillation process. In particular, applications of Puzzle often show that many attention layers can be removed with minimal accuracy loss, thus leaving sequences of FFNs uninterrupted by attention. Additionally, certain FFN layers experience significant channel pruning, leading to a reduction in their hidden dimension. Step in Figure Srgure]illustenn iilluttrates hypothetical outcome of this process.
3
FFN Fusion
Several optimization techniques reduce the number of attention layers to improve inference efficiency An attention-removed transformer block is defined by the equation
3
19
f(x) = X + (n)))
Given consecutive sequence of attention-removed blocks f f'...., fi+to we define the parallel version of this sequence f fii.tel as
fustd (X) = =X+ + 1=0
FFN (T)(
(5)
where FFN' is the FEN layer of block and 7/2 is taken from the last layer in the sequence. While in the sequential form the input for every FFN depends on the output of the previous layers, in this form all the FFN layers share the same input, which makes their computation independent and enables execution in parallel across different GPUs. Another useful property of this formulation is that equation 5 is equivalent to equation 4 with single wider FFN when the weights are concatenated.
Theorem 3.1. Let n E N, and let FFN!.... "PNN"...,FFF.... FFFN be a sequence of FEN functions, where the weights of FFN are Wi, W W3. Then, the sum of these FFN functions (equation 5) is equivalent to single FFN function FFN", with the weight matrices given by:
W = [WW)] .w.., (WT)))
W2 = [WW)] .....("!)!'
W3 [WW)).... (W)))
where [.......] denotes the concatenation of matrices along the second axis and the dimensions of the matrices are W W Bnd, xd. W3 Rddxndn
The theorem (and proof in Appendix A are written for the simple case where d, is equal for all the FFNs. The extension for the case where the hidden dimension is different for each layer is straightforward. While written in the context of Equation the theorem holds for other common FFN variants 337 39]. See visualization in Figure 7
Efficiency motivation and analysis. LLMs are ubiquitously designed in sequential blocks, with the block sizes and the number of blocks increasing as the models grow in size. For bigger models, parallelization techniques such as tensor parallel (TP) are utilized to split work across GPUs and reach acceptable inference latencies. However, this strategy does not result in linear latency improvements with the number of GPUs applied. The first reason, is the communication time required for the all-reduce that follows each TP block. The second reason is more subtle: GPUs are at their best for very large operations. As each atomic General Matrix Multiplication (GEMM) operation becomes smaller, low level overheads (GPU wave quantization) become more apparent and take a larger portion of the latency budget. Therefore, increasing the computation per GPU (by increasing the size of block) while reducing the number of synchronizations needed (by using fewer blocks) is an effective strategy for low latencies in TP setting (see Figure 2 Drawing from this motivation, we attempt fuse many sequential blocks/FFNs in L.Mss into fewer, larger blocks/FFNs. Each single reduction in the depth of the computational graph removes one unit of time spent on synchronization, and the bigger fused blocks also enable operating at higher TP numbers with enough computation assigned to each GPU. Figure different model architectures the same TP setting.
X shows the effects described in this analysis using measurements of

---

Sync
FFN2
GPU
GPU
GPU
Sync
Sync
FFN1
Fusion[FFNI. FFN2)
G
GPU
GPU
3
G
GPU
2PU
GPU 3
Smaller block, lower utilization More sync operations
Larger block, higher utilization Less sync operations
Figure 2: FFN Fusion helps reduce latency by increasing GPU utilization and by reducing syncs.
Pairwise block dependency. is reasonable to hypothesize that not every transformer block in a sequential model is equally dependent on all its predecessors. In particular, given block may depend only on subset of the blocks that precede it. To investigate this, we perform pairwise dependency analysis between blocks. Let us define h(X) f(X)= X the con- tribution of f to X, and similarly, he as the contribution of fs. We define Tiς as the con- tribution of block j when block removed from the model. We construct dependency ma- trix M E mmxm by computing the following cosine distance:
Ncxoed Bick
Doppel itod
Figure 3: Block-wise dependency heatmap of (ltra-253B-Bas before applying FFN Fusion (log- scale). Each coordinate (i,j) encodes how much block depends on block i, measured by the co- sine distance between h (X) and (X).
Mij CosineDist(Id (oaneDiss(((x),.ee (x),
That is. M.j quantifies the dependency of block j on block i. Thus, small cosine distance in- dicates that dropping block i has little effect characteristic that can be exploited for increasing parallel computation. Conversely, a large cosine distance corresponds to strong dependency, implying that sequential processing is more critical to maintain performance. We illustrate this approach in Figure 3 by constructing a dependency matrix, M, for Mltra-25B-BBat (prior to FFN Fusion; see Section 4). This matrix visually encodes the inter- dependencies among the model's layers: darker blue hues indicate weaker dependencies--sinaall promising opportunities for FFN Fusion-whil while darker red hues denote strong dependencies that would hinder parallelization or fusion. For example, 3 shows that all layers in the model de- pend on layers and causing their entire rows be colored in dark red hues. The dark blue region, marked with a dashed square, shows a sequence of FFNs with low interdependency were selected for FEN Fusion. Due to GPU memory constraints, we are unable to fuse them all once, so we divide them into fusion sequences based on the maximum size that fits within our devices. Additionally, we applied this metric explore block parallelization (Apendix B}. aiming to parallelize general LLM blocks rather than just FFNs. We constructed the dataset that was used for this evaluation following Distillation Mix
block j. suggesting relative independence
4 Producing Large-Scale Models with FFN Fusion
In this section, we describe how FFN Fusion is applied at large scale to transform a 405BBparameter Llama-based model into our more compact Ultra-253B-Baa model. Specifically, we show how identifying and fusing long sequential feed-forward blocks reduces depth without sacrificing accuracy. We utilize lightweight refinement KD and alignment phase ensure the fused model retains or even improves upon the performance of its larger predecessor. Finally, we present comprehensive evaluation of the resulting model, demonstrating significant speedups and strong results on standard benchmarks. We first run the standard Puzzle search on Llama-405B. specifying that the derivative model must obtain a 1.5x latency speedup and fit within a single NVIDIA 8x H100 node (640 GB total), and in single B100 GPU (192 GB). This yields a -arramtere model whose overall configuration shown in Appendix D Many attention layers were removed, resulting in 50 blocks contiguously arranged without interleaved attention layers.
5

---

URr-22300-See pre-OPT
Utre-2531-BBe Base--eeee Post-CPT
Lima 3.3 708 Instract
Liima 3.1 4058 Instrcct
90
80
75
70
66
60
96
80 MMLU
MATH500
HUMANEVAL
RULER 128K
Figure 4: Comparison of Ultra-253B- Base before and after applying an additional CPT.
FFN Fusion. Next, we apply FFN Fusion 3$3 to 49 of the 50 consecutive FFN layers (See Section 5.3 for ablation on leaving the last layer). Due to per-GPU memory limits, we split the layers into four sequences [66, 73), [74, 85). [86, 100], [101, 114], each fused into a single FFN. Before fusion, the baseline model achieved 84.23 on MMLU and 8.83 on MT-Bench. Remarkably, after fusing all 49 layers-layers that, if removed entirely, would damage performance severely model still maintains similar MMLU and MT-Bench scores of 82.76 and 8.35 before any additional training.
Additional Training. To recover performance, we used KD described in [3). This involved multi- stage distillation process: 54B tokens at 8k context, followed by 5B tokens each at 16k and 32k, and finally 0.8B tokens at 128k. The KD process im- proved MMLU and MT-bench scores 85.17 and 9.10, respectively. Further optimization was explored through two methods. First, inexpensive alignment via RLHF [5). Table [ demonstrated that Ultra- 253B- Base surpassed Llama-405B' capabilities, par- ticularly on the Arena Hard benchmark. We applied only a longer continual pretraining (CPT), without alignment, following the initial KD. This involved 73B tokens at context length of 8k and another 15B tokens at a context length of 258k, and also yielded strong performance, even before instruction-basee based tuning (Figure 4). Even after removing nearly half of the model's attention layers, we still improved its long-context performance, outperforming Llama-405B on RULER-128K [22]
Table I: Comparison of Ultra-253B- B-see and its parent Llama-405B after applying FFN Fusion, KD, and alignment.
Metric
Llama-105B
Ultra-253B-
MMLU Instrict
87.54
MMLU-Pro
45
71.92
72.25
Arena Hard
20
26.56
81.92
HunanEval
85.97
86.58
MT-Bhench
50
9.06
9.19
Improvements. Ultr-253BB Base achieves a 1.71x x speedup in user latency (Table[). 35x lower per-token cost at batch size 32, and reduced memory footprint with half the attention layers and 253B parameters (down from 405B). Ultra-253B-Bas -aaee breaks the efficient frontier on single H100 node, offering
Efficiency
User latency under Tensor Parallel (TP) on a single H100 node. A higher to- kens/second value indicates lower latency for single user.
Table 2:
Model
Tokens/Second
state-of-the-art accuracy and latency under the similar hardware constraints (see Figure 5). Table 2 details the user latency (tokens/second) achieved by Llama-405B, by 253B model (with and without FFN Fusion). and by Llama-3.3-70B, all under identical tensor parallel settings on a single 8xH100 node. Notably, Ultra-253B-Bass 1.71x faster than the parent for single-user decoding. On NVIDIA H200, its rate increases to 90.05 tokens/second. With speculative decoding (27 using
Llama-405B
41.44
253B PreFusion
63.38 (1.53x)
Ultra-253B-Bas
70.92 (1.71x)
Llama-3.3-70B
94.03
as draft model with no extra training Maining-Ultrr-222 -Utra-2553 Base reaches 202 tokens/s on H200. This speculative decoding latency is averaged over all MT-Bench completions. Notably, as TP size increases, the efficiency the fused FFN layers will increase further, and thus ready to benefit from improving hardware designs, such as GB200 NVI.72 nodes. Overall, Ultra-253B-B --a--5-BBBBs demonstrates that FFN Fusion coupled with the Puzzle algorithm can greatly reduce model's depth at large scale.

---

85.5
Nemotron-
84.5
Ultra-Base Aligned
83.5
Accuracy
81.5
Nemotron-
80.5
405B-Instruct
Llama-3.1-
Ultra-Base
Llama-3,3
78.5
108-Instruct
77.5
35
B5
95
Latency (Tokens/Second)
Figure 5: Accuracy latency performance of Ultra-253B-BB Base. Latency is measured single NVIDIA H100 node with tensor parallel (TP) 8, running in FP8. The red line represents the efficient frontier, highlighting models with the best wouuaacy-t-througgye tradeoff. Accuracy (MT-Bench x 10 MMLU MMLU-Pro + Arena Hard + HumanEval)/5.
5 Additional Empirical Studies
In this section, we present empirical study to validate our approach. In Section we present results for FFN Fusion on a 70B model. Next, in Section[5.2] we examine the consequences of removing FFN layers rather than fusing them, demonstrating that these layers are vital to preserving model quality. In Section (.3 we investigate the sensitivity of certain FFNs to fusion. Section 5.4 demonstrates that the FFN Fusion phenomenon is applicable to variety of models, spanning different scales and families. Finally, in Section 5.3 we provide hypothesis as to why FFN Fusion works.
5.1
FFN Fusion in 70B scale
derivative of created using the Puzzle algorithm 33. which reduces the model 49B parameters while matching the original accuracy. The Puzzle process prunes attention in specific layers, leaving two main sequences of consecutive FFN layers: layers 42-51 (10 layers) and layers 53-70 (18 layers). We evaluate FFN Fusion at four progressively increasing levels of intensity, each reducing the number of FFN layers more than the previous:
We consider
Table 3: Evaluation of FFN Fusion. The Fusion
1lama-3.1-70B-Intteee
column indicates how many FFN layers were re- placed by the corresponding number of new layers.
Model
Fusion
MMLU MT-Bench Accuracy
Baseline
80.73
8.87
81.71
Step I
26 > 12
80.64
8.72
83.92
Step 2
26 +
80.29
8,54
82.84
Step 3
26 + 3
80.39
8.30
81.69
Step 4
26 + 2
79.98
8.25
81.21
Step 1: Fuse adjacent pairs FFNs within each sequence of consecutive FFNs.
Step 2: Merge neighboring pairs from Step to form longer fused blocks of length 4 or 5.
Step 3: Fuse the entire first sequence (layers 42-50) into a single FFN, and split the second sequence (53-69) into two fused blocks.
Step 4: Fuse the entire second sequence (53-69) as well, reducing each main FFN run to single layer.
Similarly to the Ultra-2544 no additional training when applying these fusions. Table reports each variant's performance on
we chose to exclude the last FFN in each sequence. We perform
MMLL and MT-Bench, as well as combined score Accuracy = (MMLU 10 x MTT-Bench)/? We observe that step-1 fusion reduces depth with only 1 overall accuracy drop, while step-4 fusion (collapsing each sequence to single layer) sees roughly 4% drop.

---

03
Step
Accurccy (%)
Step
FEN Fusion
FFNs removed
Removing FFNs
Baseline
Step 3 KD
20 FFNs removed
53
5,4
5.5
5.6
57
5.8
5.9
Latency (ms)
Figure 6: Accuracy vs. Latency for FFN Removal vs. Fusion.
5.2 Removing FFNS vs. FEN Fusion
An intuitive alternative to fusing FFNs is to remove the same FFN layers outright. However, as shown below, large-scale removal generally leads to significant accuracy degradation, whereas fusion preserves representational capacity by retaining all parameters in single, parallel module. To determine which FFNs drop, we rely on Puzzle's ' block-importance scores and remove the least important FFNs first. We restrict removal only to the FFN-fusion regions. For the 49B model, we observe clear advantage for fusion over removal. Figure compares the accuracy-latene trade-off for two removal levels (15 or 20 FFNs removed) versus the two fusion steps. As we gradually remove more FFNs, we decrease latency but see larger accuracy drops. By contrast, fusing step yields 10% latency improvement with only 3.5% accuracy drop. Removing 15 or 20 FFNs yields comparable or slightly higher latency gains (7.5% or 11%), but at steeper accuracy declines (5% or 8.7%). Moreover, brief Knowledge-distilaaton phase (25B tokens) on the step-3 fused model actually surpasses the baseline (MMLU: 80.56, MT-Bench: 9.00).
5.3 The Final FFN in Each Sequence is Sensitive to Fusion
With the 49B model, we observe that fusing the final FFN in long sequence often degrades accuracy more than fus- ing earlier FFNs. Below, we summarize an ablation study to highlight this phenomenon. Table Pable/4exxc exammnes various ways of fusing the two main FFN sequences (layers 42-51 and 53-70) in the 49B model. The first part of the table compares fusing the entire second sequence (53,, 69) vs. [54, 70)) and the entire first sequence ([42, 50) vs. [43, 51)) in isolation, indicating that layer 70 is especially sensitive. In the second part, we fuse both sequences simultaneously. When either layer 51 or layer 70 is included in these large fused blocks, we see additional performance drops. No- tably, skipping the final FFN (e.g. fusing [42, 50], [53, 69)) yields higher
Table 4: Evaluating the impact of fusing the final FFN in 49B scale. Fusing these final layers often causes more accuracy loss.
Sequence
MMLU
MTT-Bench
No Fusion
80.73
8.87
533 69)
80.57
8.49
154, 70
79.89
8.12
42, 50] 43, 51
80.49
8.39
80.46
8.42
42, 50),
[53, 69
79.98
8.25
scores than including layer 51 sequence
42,
53. 69
80.05
7.64
51
54. 70
79.92
7.38
70. The final FFN in each attention-remors appears uniquely important to the model's representations. While most layers can be safely fused, incorporating the last FFN often triggers a significant accuracy drop. Consequently, omitting this final FFN from the fused groups typically a more reliable choice for efficient fusion with minimal performance loss.
|2, 51],
[33, 70
79.89
7.30

---

5,4 Additional Models
This section explores FFN Fusion on additional mod- els, evaluating its effectiveness across different scales and model families. We applied FFN Fusion to
Table 5: FFN Fusion impact on Mistral Large 2 and Llama 3.1 8B Instruct.
1lama-3.1-88B I-stuunt 13 and Mistral Large 2 follow- ing a similar procedure as with our 253B and 49B models. Specifically, we used the Puzzle algorithm on both models to remove a portion of their attention layers. Following this, we fused consecutive sequences of FFNs. For Mistral Large 2. we removed 27 attention layers that resulted 2 sequences of 14 and 13 consecutive FEN layers. We fused each sequence (excluding the last FFN) to a single FFN layer. For Llama-3.1-88
Fusion
MMLU MT-Bench
Mistral Large
80.56
7.16
80.06
6.76
Llama 3.1 8B Instruct
68.91
8.28
68.22
7.99
(32 layers), 8 consecutive attention layers. We fused the sequence (excluding the last FFN) to single FFN layer. We also conducted short KD training on the 8B model over 20B tokens and compared the results to the original model. Table
removed
5 provides a clear demonstration of FFN Fusion's generalizability, showing its effectiveness across additional models.
5.5 Fusion Explainability
In this section, we aim further explain the phenomenon of FFN Fusion. We examine the functional structure of LLMs that allows for the fusing of layers, specifically focusing on the consecutive FFN layers after attention removal. First, we rewrite the token-level normalization equation in an alternative form: (1;) = De,," xi/ |||/]|, E Rdd
Table 6: Reverse order vs Fusion experi- ment on Puzzle-49B.
Fusion Accuracy
Reverse
Sequence
Accuracy
42, 49
83.74
83.52
where we replace the normalization parameters with diagonal matrix D E Rdxdd with 8% on the diago- nal. This reformulation allows us to "push" the matrix into its corresponding module- eittee attention or FFN and consider equations I and 4 functions of the token direction /7 = = while ignoring the mag-
53, 60
83.24
83.25
63, 67
84.22
84.10
66,, 70
84.00
83.88
nitude. Figure (a and b), Similar to the measurements in (30]. shows the cosine distance between X and f(X) for each layer of both models. The plot reveals that the FFN fusing areas-[42, (42,5]] [53, 70] lower cosine distance compared to other regions in the model. Assuming that small changes in the input to FFN layer lead to small changes in its output, we can conclude that fusing the FFNs with low cosine distance (except for the last one) will not cause significant changes to the model. Specifically, by fusing the FFN layers, we are altering the input for each layer, but the directional difference between the original input (according to the original model's structure) and the new input (resulting from fusion) remains small, and the outputs should remain similar under the smoothness assumption. Another experiment (Table 6) that strengthens this claim shows that even if we reverse the FFNs order instead of fusing them, we get similar performances. The results in Figure 11 (c and d) present the relationship between the layer input X and h(X)= f(X) X. The small change in token directions may be linked to the low cosine distance between h(X) and X, but they are nearly orthogonal (with a distance of around 0.95 for both models across all layers). However, more suitable explanation comes from the ratio [] X) II ||X]] []. which is smaller in the fused regions. This suggests that Á (X) is small compared X that it does not significantly affect the direction of X. Although this section can also serves as motivation for FFN removal, we demonstrated in 6 removing these layers significantly harms the model's 'ss performance, much greater extent than fusion does. Both metries were evaluated using the Distillation Mix dataset.
for the 49B model and [66, 115] for the 253B model-exhine
Concluding Remarks
This work introduces FFN Fusion, a novel optimization technique that demonstrates how sequential computation in large language models can be effectively parallelized. Through comprehensive experiments at scales from 49B to 253B parameters, we showed that sequences of FFN layers can be fused with minimal impact on model capabilities. Our work revealed couple of key insights: (1) FFN layers in attention-remover regions exhibit remarkably low inter-layer dependencies, suggesting these

---

computations may be more independent than the sequential architecture implies; and (2) empirically. final FFN layers in parallelizable sequences show higher sensitivity to fusion-ar an observation that points to potentially meaningful structural patterns in how these models process information.
The practical impact of our work is demonstrated through Ultra-553B- which achieves 1.71x speedup in inference latency compared to its parent Llama-405B model, while requiring 35x lower per-token cost at batch size 32. Most remarkably, these efficiency gains come with minimal degradation in model capability-in fact, Ultra-253B-a B-Base matches or exceeds its parent's performance across key benchmarks despite using only 253B parameters compared to the original 405B.
Our findings open several promising research directions. First, the clear patterns in inter-layer dependencies we observed could provide a new lens for model interpretability research. The ability to identify which FFN layers operate independently versus those that require sequential processing may offer insights into how these models structure and transform their internal representations. Second, our preliminary finding that even full transformer blocks can sometimes be parallelized suggests possibilities for new architectural designs explicitly optimized for parallel execution. Third, future work could extend FFN Fusion to models incorporating Mixture-of-Expere (MoE) layers. Finding a way to fuse MoE layers while maintaining sparse activation patterns and efficient expert routing could lead efficiency improvements. Finally. the strong performance of FFN Fusion at larger scales raises intriguing questions about the relationship between model size and natural parallelization.
7 Societal Impact
The inference efficiency achieved through FFN Fusion, as demonstrated by its impact on operational costs in Section 4 significantly improves the accessibility of large language models. By reducing these costs, a wider array of users and organizations can now better afford to deploy and leverage these powerful technologies, effectively dismantling economic barriers.
References
[1] M. R. A14Science and M. A. Quantum. The impact of large language models on scientific preliminary study using gpt-4, 2023.
discovery:
[2] R. Anil, S. Borgeaud, Y. Wu, J. Alayrac, J. Yu, R. Soricut, J. Schalkwyk, A. M. Dai, A. Hauth,
K. Millican, D. Silver, S. Petrov, M. Johnson, I. Antonoglou, J. Schrittwieser, A. Glaese, J. Chen,
E. Pitler, T. P. Lillicrap, A. Lazaridou, O. Firat, J. Molloy, M. Isard, P. R. Barham, T. Hennigan, B. Lee, F. Viola, M. Reynolds, Y. Xu, R. Doherty, E. Collins, C. Meyer, E. Rutherford, E. Moreira, K. Ayoub, M. Goel, G. Tucker, E. Piqueras, M. Krikun, I. Barr, N. Savinov,
I. Danihelka, B. Roelofs, A. White, A. Andreassen, T. von Glehn, L.. Yagati, M. Kazemi, L. Gonzalez, M. Khalman, J. Sygnowski, and et al. Gemini: A family of highly capable multimodal models. CoRR, abs/23 12.12322.1180., 2023.
[3] A. Bercovich, T. Ronen, T. Abramovich, N. Ailon, N. Assaf, M. Dabbah, I. Galil, A. Geifman, Y. Geiffman, I. Golan, N. Haber, E. Karpas, R. Koren, I. Levy, Molchanov, Mor, Z. Moshe, N. Nabwani, O. Puny. R. Rubin, I. Schen, I. Shahaf, O. Tropp, O. U. Argov, R. Zilberstein, and
R. Eaniv.... Puzzle: Distillation-based nas for inference-optimizee 11ms, 2024.
[4] M. Chen, J. Tworek, H. Jun, Q. Yuan, H. P. de Oliveira Pinto, J. Kaplan, H. Edwards, Y. Burda, N. Joseph, G. Brockman, A. Ray, R. Puri, G. Krueger, M. Petrov, H. Khlaaf, G. Sastry, P. Mishkin, B. Chan, S. Gray, N. Ryder, M. Pavlov, A. Power, L. Kaiser, M. Bavarian, C. Winter, Tillet, F. P. Such, D. Cummings, M. Plappert, F. Chantzis, Barnes, A. Herbert- Voss, W. H. Guss, A. Nichol, A. Paino, N. Tezak, J. Tang, I. Babuschkin, S. Balaji, S. Jain, W. Saunders, C. Hesse, A. N. Carr, J. Leike, Achiam, V. Misra, E. Morikawa, A. Radford, M. Knight, M. Brundage, M. Murati, K. Mayer, P. Welinder, B. McGrew, D. Amodei, S. McCandlish, Sutskever, and W. Zaremba. Evaluating large language models trained on code. CORR, abs/2107.03374 2021.
[5] DeepSeek-AI, D. Guo, D. Yang, H. Zhang, Song, R. Zhang, R. Xu, Q. 7hu, Ma, Wang, X. Bi, X. Zhang, X. Yu, Y. Wu, Z. F. Wu, Z. Gou, Z. Shao, Z. Li, Z. Gao, A. Liu, B. Xue,
10

---

B. Wang. B. Wu, B. Feng, C. Lu, C. Zhao, C. Deng, C. Zhang, C. Ruan, D. Dai, D. Chen,
D. Ji, E. Li, F. Lin, F. Dai, F. Luo, G. Hao, G. Chen, G. Li, H. Zhang. H. Bao, H. Xu, H. Wang. H. Ding, H. Xin, H. Gao, H. Qu, H. Li, J. Guo, J. J. Wang, J. Chen, J. Yuan, J. Qiu, J. Li,
J. Cai, Ni, Liang. J. Chen, K. Dong, K. Hu, K. Gao, K. Guan, K. Huang, K. Yu, L. Wang.
I.. Zhang. L.. Zhao, L.. Wang. I. Zhang, L. Xu, L. Xia, M. Zhang, M. Zhang, M. Tang, M. Li,
M. Wang, M. Li, N. Tian, P. Huang, P. Zhang, Q. Wang. Q. Chen, Q. Du, R. Ge, R. Zhang,
R. Pan, R. Wang. R. J. Chen, R. L. Jin, R. Chen. S. Lu, S. Zhou, S. Chen, S. Ye, S. Wang. S. Yu,
S. Zhou, S. Pan, S. S. Li, S. Zhou, S. Wu, S. Ye, T. Yun, T. Pei, T. Sun, T. Wang, W. Zeng,
W. Zhao, W. Liu, W. Liang, W. Gao, W. Yu, W. Zhang. W. L. Xiao, W. An, X. Liu, X. Wang,
X. Chen, X. Nie, X. Cheng, X. Liu, X. Xie, X. Liu, X. Yang, X. Li, X. Su, X. Lin, X. Q. Li,
X. Jin, X. Shen, X. Chen, X. Sun, X. Wang, X. Song, X. Zhou, X. Wang. X. Shan, Y. K. Li, Q.Q. Wang. Y. X. Wei, Y. Zhang, Y. Xu, Y. Li, Y. Zhao, Y. Sun, Y. Wang, Y. Yu, Y. Zhang, Y. Shi,
Y. Xioog, Y. He, Y. Piao, Y. Wang. Y. Tan, Y. Ma, Y. Liu, Y. Guo, Y. Ou, Y. Wang, Y. Gong,
Y. Zou, Y. He, Y. Xiong. Y. Luo, Y. You, Y. Liu, Y. Zhou, Y. X. Zhu, Y. Xu, Y. Huang, Y. Li,
Y.Zheng, Y. Zhu, Y. Ma, Y. Tang. Y. Zha, Y. Yan, Z.Z. Ren, Z. Ren, Z. Sha, Z. Fu, Z. Xu, Z. Xie,
Z. Zhang, Z. Hao, Z. Ma. Z. Yan, Z. Wu, Z. Gu, Z. Zhu, Z. Liu, Z. Li, Z. Xie, Z. Song, Z. Pan,
Z. Huang, Z. Xu, Z. Zhang, and Z. Zhang. Deepseek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning, 2025.
[6] DeepSeek-AI, A. Liu, B. Feng, B. Xue, B. Wang, B. Wu, C. Lu, C. Chao C. Deng, C. Zhang, C. Ruan, D. Dai, D. Guo, D. Yang. D. Chen, D. Ji, E. Li, F. Lin, F. Dai, F. Luo, G. Hao, G. Chen,
G. Li, H. Zhang. H. Bao, H. Xu, H. Wang. H. Zhang, H. Ding, H. Xin, H. Gao, H. Li, H. Qu, J. L. Cai, J. Liang, J. Guo, J. Ni, J. Li, Wang, Chen, J. Chen, J. Yuan, J. Qiu, J. Li, J. Song, K. Dong, K. Hu, K. Gao, K. Guan, K. Huang. K. Yu, L. Wang, L. Zhang, L. Xu, L. Xia, L. Zhao,
L. Wang, L. Zhang, M. Li, M. Wang. M. Zhang, M. Zhang, M. Tang, M. Li, N. Tian, Huang.
P. Wang. P. Zhang, Q. Wang. Q. Zhu, Q. Chen. Q. Du, R. J. Chen, R. L. Jin, R. Ge. R. Zhang, R. Pan, R. Wang, R. Xu, R. Zhang, R. Chen, S. Li, Lu, S. Zhou, S. Chen, Wu, S. Ye, S. Ye, S. Ma, S. Wang, S. Zhou, S. Yu, S. Zhou, S. Pan, T. Wang. T. Yun, T. Pei, T. Sun, W. L.. Xiao, and W. Zeng. Deepseek-v3 technical report. CORR, abs/2412. 194377, 2024.
[7] T. Dettmers, M. Lewis, Y. Belkada, and L. Zettlemoyer. Gpt3. int8 (): 8-bit matrix multiplication for transformers at scale. Advances in neural information processing systems, 05:30318-30332, 2022.
[8] T. Dettmers and L. Zettlemoyer. The case for 4-bit precision: k-bit inference scaling laws. In International Conference on Machine Learning, pages 7750-7774. PMLR, 2023.
[9] A. Dubey, A. Jauhri, A. Pandey, A. Kadian, A. Al-Dahle, A. Letman, A. Mathur, A. Schelten, A. Yang. A. Fan, A. Goyal, A. Hartshorn, A. Yang, A. Mitra, A. Sravankumar, A. Korenev, A. Hinsvark, A. Rao, A. Zhang, A. Rodriguez, A. Gregerson, A. Spataru, B. Rozière, B. Biron, B. Tang, B. Chern, C. Caucheteux, C. Nayak, C. Bi, C. Marra, C. McConnell, C. Keller,
C. Touret, C. Wu, C. Wong, C. C. Ferrer, C. Nikolaidis, D. Allonsius, D. Song, D. Pintz, D. Livshits, D. Esiobu, D. Choudhary, D. Mahajan, D. Garcia-Olano, D. Perino, D. Hupkes, E. Lakomkin, E. AlBadawy, E. Lobanova, E. Dinan, E. M. Smith, F. Radenovic, F. Zhang, G. Synnaeve, G. Lee, G. L. Anderson, G. Nail, G. Mialon, G. Pang, G. Cucurell, H. Nguyen, H. Korevaar, H. Xu, H. Touvron, I. Zarov, I. A. Ibarra, I. M. Kloumann, I. Misra, I. Evtimov, J. Copet, J. Lee, J. Geffert, J. Vranes, J. Park, J. Mahadcokar, J. Shah, J. van der Linde, J. Billock, J. Hong, J. Lee, J. Fu, J. Chi, J. Huang. J. Liu, J. Wang. J. Yu, J. Bitton, J. Spisak, J. Park, J. Rocca, J. Johnstun, J. Saxe, J. Jia, K. V. Alwala, K. Upasani, K. Plawiak, K. Li, K. Heafield, K. Stone, and et al. The llama 3 herd of models. CoRR, abs/2407.22 .27833, 2024.
[10] A. Fan, E. Grave, and A. Joulin. Reducing transformer depth on demand with structured dropout. arXiv preprint arXiv: 1909. 11556, 2019.
[11] J. Frankle and M. Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural 03635, 2018
networks. arXiv preprint
[12] E. Frantar, S. Ashkboos, T. Hoefter, and D. Alistarh. Gptq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprint arXiv:2210.17323, 2022.
[13] A. Grattafiori, A. Dubey, A. Jauhri, A. Pandey, A. Kadian, A. Al-Dahle, A. Letman, A. Mathur, A. Schelten, A. Vaughan, A. Yang, A. Fan, A. Goyal, A. Hartshorn, A. Yang, A. Mitra,
11

---

A. Sravankumar, A. Korenev, A. Hinsvark, A. Rao, A. Zhang, A. Rodriguez, A. Gregerson,
A. Spataru, B. Roziere, B. Biron, B. Tang, B. Chemn C. Caucheteux, C. Nayak, C. Bi, C. Marra, C. C.Moonnell, C. Keller, C. Touret, C. Wu, C. Wong, C. C. Ferrer, C. Nikolaidis, D. Allonsius, D. Song. D. Pintz, D. Livshits, D. Wyatt, D. Esiobu, D. Choudhary, D. Mahajan, D. Garcia- Olano, D. Perino, D. Hupkes, E. Lakomkin, E. AlBadawy, E. Lobanova, E. Dinan, E. M. Smith, F. Radenovic, F. Guzmán, F. Zhang, G. Synnaeve, G. Lee, G. L. Anderson, G. Thattai, G. Nail, G. Mialon, G. Pang. G. Cucurell, H. Nguyen, H. Korevaar, H. Xu, H. Touvron, I. Zarov, J. A. Ibarra, I. Kloumann, I. Misra, I. Evtimov, J. Zhang. J. Copet. J. Lee, J. Geffert, J. Vranes, J. Park, J. Mahadeokar, J. Shah, J. van der Linde, J. Billock, J. Hong. J. Lee, J. Fu, J. Chi, J. Huang, J. Liu, J. Wang, J. Yu, J. Bitton, J. Spisak, J. Park, J. Rocca, J. Johnstun, J. Saxe. J. Jia. K. V. Alwala, K. Prasad, K. Upasani, K. Plawiak, K. Li, K. Heafield, K. Stone, K. El-Arini, K. Iyer, K. Malik, K. Chiu, K. Bhalla, K. Lakhotia, I.. Mantala-Yeary, L. van der Maaten, L. Chen, L. Tan, I. Jenkins, I.. Martin, L.. Madaan, I.. Malo, I.. Blecher, L. Landzaat, I.. de Oliveira, M. Muzzi, M. Pasupuleti, M. Singh, M. Paluri, M. Kardas, M. Tsimpoukelli, M. Oldham, M. Rita, M. Pavlova, M. Kambadur, M. Lewis, M. Si, M. K. Singh, M. Hassan, N. Goyal, N. Torabi, N. Bashlykov, N. Bogoychev, N. Chatterji, N. Zhang, O. Duchenne, O. Çelebi, Alrassy, P. Zhang, P. Li, P. Vasic, P. Weng, P. Bhargava, P. Dubal, P. Krishnan, S. Koura, P. Xu, Q. He, Q. Dong, R. Srinivasan, R. Ganapathy, R. Calderer, R. S. Cabral, R. Stojnic, R. Raileanu, R. Maheswari, R. Girdhar, R. Patel, R. Sauvestre, R. Polidoro, R. Sumbaly, R. Taylor, R. Silva, R. Hou, R. Wang, S. Hosseini, S. Chennabasappa. S. Singh, S. Bell, S. S. Kim. S. Edunov, S. Nie, S. Narang, S. Ruparthy, S. Shen, S. Wan. S. Bhosale, S. Zhang. S. Vandenhende, S. Batra, S. Whitman, S. Sootla, S. Collot, S. Gururangan, S. Borodinsky, Herman, T. Fowler, T. Sheasha, T. Georgiou, T. Scialom, T. Speckbacher, T. Mihaylov, T. Xiao, U. Karn, V. Goswami, V. Gupta, V. Ramanathan, V. Kerkez, V. Gonguet, V. Do, V. Vogeti, V. Albiero, V. Petrovic, W. Chu, W. Xiong, W. Fu, W. Meers, X. Martinet, X. Wang. X. Wang, X. E. Tan, X. Xia, X. Xie, X. Jia, X. Wang. Y. Goldschlag, Y. Gaur, Y. Babaei, Y. Wen, Y. Song, Y. Zhang, Y. Li, Y. Mao, Z. D. Coudert, Z. Yan, Z. Chen, Z. Papakipos, A. Singh, A. Srivastava, A. Jain, A. Kelsey, A. Shajnfeld, A. Gangidi, A. Victoria, A. Goldstand, A. Menon, A. Sharma, A. Boesenberg. A. Baevski, A. Feinstein, A. Kallet, A. Sangani, A. Teo, A. Yunus, A. Lupu, A. Alvarado, A. Caples, A. Gu, A. Ho, A. Poulton, A. Ryan, A. Ramchandani, A. Dong, A. Franco, A. Goyal, A. Saraf, A. Chowdhury, A. Gabriel, A. Bharambe, A. Eisenman, A. Yazdan, B. James, B. Maurer, B. Leonhardi, B. Huang, B. Loyd, B. D. Paola, B. Paranjape, B. Liu, B. Wu, B. Ni, B. Hancock, B. Wasti, B. Spence, B. Stojkovic, B. Gamido, B. Montalvo, C. Parker, C. Burton, C. Mejia, C. Liu, C. Wang, C. Kim, C. Zhou, C. Hu, C.-H. Chu, C. Cai, C. Tindal, C. Feichtenhofer, C. Gao, D. Civin, D. Beaty, D. Kreymer, D. Li, D. Adkins, D. Xu, D. Testuggine, D. David, D. Parikh, D. Liskovich, D. Foss, D. Wang, D. Le. D. Holland, E. Dowling, E. Jamil, E. Montgomery, E. Presani, E. Hahn, E. Wood, E.-T. Le, E. Brinkman, E. Arcaute, E. Dunbar, E. Smothers, F. Sun, F. Kreuk, F. Tian, F. Kokkinos, F. Ozgenel, F. Caggioni, F. Kanayet, F. Seide, G. M. Florez, G. Schwarz, G. Badeer, G. Swee, G. Halpern, G. Herman, G. Sizov. Guangyi. Zhang, G. Lakshminarayanann H. Inan, H. Shojanazeri, H. Zou, H. Wang, H. Zha, H. Habeeb, H. Rudolph, H. Suk, H. Aspegren, H. Goldman, H. Zhan, I. Damlaj, I. Molybog, I. Tufanov, I. Leontiadis, I.-F. Veliche, I. Gat, J. Weissman, J. Geboski, J. Kohli, J. Lam, J. Asher, J.-B. Gaya, J. Marcus, J. Tang, J. Chan, J. Zhen, J. Reizenstein, J. Teboul, J. Zhong. J. Jin, Yang. J. Cummings, J. Carvill, J. Shepard, J. McPhie, J. Torres, J. Ginsburg, J. Wang, K. Wu, K. H. U, K. Saxena, K. Khandelwal, K. Zand, K. Matosich, K. Veeraraghavan, K. Michelena, K. Li, K. Jagadeesh, K. Huang. K. Chawla, K. Huang, L. Chen, L. Garg, L. A, L. Silva, L. Bell, L. Zhang, L. Guo, L. Yu, L. Moshkovich, L. Wehrstedt, M. Khabsa, M. Avalani, M. Bhatt, M. Mankus, M. Hasson, M. Lennie, M. Reso, M. Groshev, M. Naumov, M. Lathi, M. Keneally, M. Liu, M. L. Seltzer. M. Valko, M. Restrepo, M. Patel, M. Vyatskov, M. Samvelyan, M. Clark, M. Maccy, M. Wang. M. J. Hermoso, M. Metanat, M. Rastegari, M. Bansal, N. Santhanam, N. Parks, N. White, N. Bawa, N. Singhal, N. Egebo, N. Usunier, N. Mehta, N. P. Laptev, N. Dong, N. Cheng, O. Chernoguz, O. Hart, O. Salpekar, O. Kalinli, P. Kent, P. Parekh, P. Saab, P. Balaji, P. Rittner, P. Bontrager, P. Roux, P. Dollar, P. Zvyagina, P. Ratanchandani, P. Yuvraj, Q. Liang, R. Alao, R. Rodriguez, R. Ayub, R. Murthy, R. Nayani, R. Mitra, R. Parthasarathy, R. Li, R. Hogan, R. Battey, R. Wang, R. Howes, R. Rinott, S. Mehta, S. Siby, S. J. Bondu, S. Datta, S. Chugh, S. Hunt, S. Dhillon, S. Sidorov, S. Pan, S. Mahajan, S. Verma, S. Yamamoto, S. Ramaswamy, S. Lindsay, S. Lindsay, S. Feng, S. Lin,
S. C. Zha, S. Patil, S. Shankar, S. Zhang, S. Zhang, S. Wang, S. Agarwal, S. Sajuyighe,
S. Chintala, S. Max. S. Chen, S. Kehoe, S. Satterfield, S. Govindaprasad, S. Gupta, S. Deng,
12

---

S. Cho. S. Virk, S. Subramanian, S. Choudhury, S. Goldman, T. Remez, T. Glaser, T. Best, T. Koehler, T. Robinson, T. Li. T. Zhang. T. Matthews, Chou, T. Shaked, V. Vontimitta, V. Ajayi, V. Montanez, V. Mohan, V. S. Kumar, V. Mangla, V. Jonescu, V. Poenaru, V. T. Mihailescu, V. Ivanov, W. Li, W. Wang, W. Jiang, W. Bouaziz, W. Constable, X. Tang, X. Wu, X. Wang, X. Wu, X. Gao, Y. Kleinman, Y. Chen, Y. Hu, Y. Jia, Y. Qi, Y. Li, Y. Zhang, Zhang. Y. Adi, Y. Nam, Yu, Wang. Y. Zhao, Y. Hao, Y. Qian, Y. Li, Y. He, Z. Rait, Z. DeVito, Z. Rosnbrick, Z. Wen, Z. Yang, Z. Zhao, and Z. Ma. The llama 3 herd of models, 2024.
[14] A. Gromov, K. Tirumala, H. Shapourian, P. Glorioso, and D. A. Roberts. The unreasonable ineffectiveness of the deeper layers, 2024.
[15] X. Guan, L. L. Zhang, Y. Liu, N. Shang, Y. Sun, Y. Zhu, Yang, and M. Yang. rstar-math: Small Iims can master math reasoning with self-evolved deep thinking, 2025.
[16] Y. Guan, D. Wang, Z. Chu, S. Wang, F. Ni, R. Song, L. Li, J. Gu, and C. Zhuang. Intelligent virtual assistants with Ilm-based process automation, 2023.
[17] S. Han, J. Pool, J. Tran, and W. Dally. Learning both weights and connections for efficient neural network. Advances in neural information processing systems, 28, 2015.
[18] B. Hassibi and D. Stork. Second order derivatives for network pruning: Optimal brain surgeon. Advances in neural information processing systems, 5, 1992.
[19] S. He, G. Sun, Z. Shen, and A. Li. What matters in transformers? not all attention is needed, 2024.
[20] D. Hendryeks and K. Gimpel. Bridging nonlinearities and stochastic regularizers with gaussian error linear units. CoRR,
2016.
[21]
LeeeDiigtaa-Teechollggii
[22] C.-P. Hsich, S. Sun, S. Kriman, S. Acharya, D. Rekesh, F. Jia, Y. Zhang, and B. Ginsburg. Ruler: What's the real context size of your long-context language models? arXiv preprint arXiv:2404.0666 066554 2024.
[23] A. Q. Jiang, A. Sablayrolles, A. Roux, A. Mensch, B. Savary, C. Bamford, D. S. Chaplot, D. de las Casas, E. B. Hanna, F. Bressand, G. Lengyel, G. Bour, G. Lample, L. R. Lavaud,
L. Saulnier, M.-A. Lachaux, P. Stock, S. Subramanian, S. Yang, S. Antoniak, T. L. Scao, Gervet, T. Lavril, T. Wang. T. Lacroix, and W. E. Sayed. Mixtral of experts. 2024.
[24] W. Kwon, Z. Li, S. Zhuang, Sheng, L. Zheng, C. H. Yu, J. E. Gonzalez, H. Zhang, and I. Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles, 2023.
[25] V. Lad, W. Gurnee, and M. Tegmark. The remarkable robustness of Ilms: Stages of inference?, 2024.
[26] Y. LeCun, J. Denker, and S. Solla. Optimal brain damage. Advances in neural information processing systems, 2, 1989.
[27] Y. Leviathan, M. Kalman, and Y. Matias. Fast inference from transformers via speculative decoding. In A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett, editors, International Conference on Machine Learning. ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pages 19274-19286. PMLR, 2023.
[28] H. Li, A. Kadav, I. Durdanovic, H. Samet, and H. P. Graf. Pruning filters for efficient convnets. arXiv preprint
2016.
[29] T. Li, W. Chiang, E. Frick, L. Dunlap, T. Wu, B. Zhu, J. E. Gonzalez, and Stoica. From crowdsourced data to high-quality benchmarks; Arena-hard and benchbuilder pipeline. CoRR, abs/2406.1 1939, 2024.
[30] 7. Liu, J. Wang, T. Dao, T. Zhou, B. Yuan, Z. Song, A. Shrivastava, C. Zhang, Y. Tian, C. Re, and B. Chen. Deja vu: Contextual sparsity for efficient 1lms at inference time, 2023.
13

---

[31] P. Michel, O. Levy, and G. Neubig. Are sixteen heads really better than one? Advances in neural information processing systems. 32, 2019.
[32] Nvidia, ;, B. Adler, N. Agarwal, A. Aithal, D. H. Anh, P. Bhattacharya, A. Brundyn, J. Casper, B. Catanzaro, S. Clay, J. Cohen, S. Das, A. Dattagupta, O. Delalleau, L. Derezynski, Y. Dong, D. Egert, E. Evans, A. Ficek, D. Fridman, S. Ghosh, B. Ginsburg. I. Gitman, T. Grzegorzek, R. Hero, J. Huang, V. Jawa, J. Jennings, A. Jhunjhunwala, J. Kamalu, S. Khan, O. Kuchaiev,
P. LeGresley, H. Li, J. Liu, 7. Liu, F. Long, A. S. Mahahaleshwarkar, S. Majumdar, J. Maki, M. Martinez, M. R. de Melo, I. Moshkov, D. Narayanan, S. Narenthiran, J. Navarro, P. Nguyen, O. Nitski, V. Noroozi, G. Nutheti, C. Parisien, J. Parmar, M. Patwary, K. Pawelec, W. Ping, Prabhumoye, R. Roy. T. Saar, V. N. Sabavat, S. Satheesh, J. P. Scoweroft, J. Sewall, P. Shamis, G. Shen, M. Shoeybi, D. Sizer, M. Smelyanskiy, F. Soares, M. N. Sreedhar, D. Su, Subramanian, S. Sun, S. Toshniwal, H. Wang, Z. Wang, You, J. Zeng. J. Zhang, J. Zhang, V. Zhang, Y. Zhang, and C. Zhu. Nemotron-4 340b technical report, 2024.
[33] OpenAl. GPT-4 technical report. CoRR, abs/2303.08774, 2023.
[34] G. Penedo, H. Kydlícek, L.. B. Allal, A. Lozhkov, M. Mitchell, C. Raffel, L.. von Werra, and Wolf. The fineweb datasets: Decanting the web for the finest text data at scale. CoRR, abs/2406. 17557, 2024.
[35] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever. Language models are unsupervised multitask learners. OpenAl, 2019. Accessed: 2024-11-15.
[36] H. Sajjad, F. Dalvi, N. Durrani, and P. Nakov. On the effect of dropping layers of pre-trained transformer models. Computer Speech & Language, 77:101429, 2023.
[37] N. Shazeer. Glu variants improve transformer, 2020.
[38] N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. V. Le, G. E. Hinton, and J. Dean. Outra- geously large neural networks: The sparsely-gated mixture-of-experts layer. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017.
[39] D. R. So, W. Mańke, H. Liu, Z. Dai, N. Shazeer, and Q. V. Le. Primer: Searching for efficient transformers for language modeling, 2022.
[40] L. Soldaini, R. Kinney, A. Bhagia, D. Schwenk, D. Atkinson, R. Authur, B. Bogin, K. R. Chandu, J. Dumas, Y. Elazar, V. Hofmmann, A. H. Jha, S. Kumar, L. Lucy, X. Lyu, N. Lambert,
Magnusson, Morrison, N. Muennighoff, A. Naik, C. Nam, M. E. Peters, A. Ravichander, K. Richardson, Z. Shen, E. Strubell, N. Subramani, O. Tafjord, E. P. Walsh, L. Zettlemoyer, N. A. Smith, H. Hajishirzi, Beltagy, D. Groeneveld, Dodge, and K. Lo. Dolma: an open corpus of three trillion tokens for language model pretraining research. In L. Ku, A. Martins, and V. Srikumar, editors, Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2024, Bangkok, Thailand, August 11-16, 2024, pages 15725-15788. Association for Computational Linguistics, 2024.
[41] R. Taylor, M. Kardas, G. Cucurull, T. Scialom, A. Hartshorn, E. Saravia, A. Poulton, V. Kerkez, and R. Stojnic. Galactica: A large language model for science, 2022.
[42] A. Vaswani, N. Shazeer, N. Parmar, Uszkoreit, L. Jones, A. N. Gomez, Kaiser, and
I. Polosukhin. Attention is all you need, 2023.
[43] E. Voita, D. Talbot, F. Moiseev, R. Sennrich, and I. Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. arXiv preprint arXiv:1905.099 094188 2019.
[44] Y. Wang, X. Ma. G. Zhang, Y. Ni, A. Chandra, S. Guo, W. Ren, A. Arulraj. X. He, Z. Jiang, T. Li, M. Ku, K. Wang, A. Zhuang, R. Fan, X. Yue, and W. Chen. Mmlu-pro: A more robust and challenging multi-task language understanding benchmark. CoRR, 0as/2400.01574. 2024.
[45] Z. Wang, A. Bukharin, O. Delalleau, D. Egert, G. Shen, J. Zeng, O. Kuchaiev, and Y. Dong. Hedpsteer2-pprefeeence Complementing ratings with preferences. CoRR, abs/2410.01257, 2024.
14

---

[46] T. Wolf, L. Debut, V. Sanh, J. Chaumond, C. Delangue, A. Moi, P. Cistac, T. Rault, R. Louf, M. Funtowiez, Davison, Shleifer, von Platen, C. Ma. Y. Jemite, J. Plu, C. Xu, T. Scao,
S. Gugger, M. Drame, Q. Lhoest, and A. M. Rush. Huggingface's transformers: State-of-the-art natural language processing, 2020.
[47] G. Xiao, J. Lin, M. Seznec, H. Wu, J. Demouth, and Han. Smoothquant: Accurate and efficient post-training quantization for large language models. In International Conference on Machine Learning, pages 38087-38099. PMLR, 2023.
[48] Z. Xu, F. Jiang, L. Niu, Y. Deng, R. Poovendran, Y. Choi, and B. Y. Lin. Magpie: Aligmment data synthesis from scratch by prompting aligned Iims with nothing, CoRR, abs/2406.08 08464, 2024.
[49] M. Zhang and He. Accelerating training of transformer- based language models with pro- gressive layer dropping. Advances in neural information processing systems, 33:14011-14023, 2020.
[50] I.. Zheng, Chiang, Sheng, S. Zhuang. 7.. Wu, Y. Zhuang. Lin, Z. Li, D. Li, E. P. Xing, H. Zhang, J. E. Gonzalez, and I. Stoica. Judging llm-as-a-judge with mt-bench and chatbot arena. In A. Oh, T. Naumann, editors, Advances in Neural Information Processing Systems 36: Annual Conference Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 16, 2023, 2023.
Globerson, K. Saenko, M. Hardt, and Levine,
A Proof of Theorem 3.1
Proof. The proof will proceed by induction. For the base case, we consider two FFN layers, FFN1 and FFN2, We will demonstrate that FFN" FFN' + +FFN2, The induction step follows naturally from the associative property of addition.
(FFN' (X) + FFN? (X)) =(o(((W)))) @ XXVI))(33)++((XXX333 @ ((W))))))))()))))) (WW3)
= =(o(X(W3))) O ((())) x(W?)) (Wg)*
( o [X(W))* x(W3)
=(O(X(WG)") R (WW)) =XFN'(X)
where [, is the concatenation of matrices along the second dimension.
B Block Parallelization
In all previous experiments, we focused fusing sequences of FFNs in attention- removed layers (see Section 33). In that setting, the homogeneous structure of FFN-only layers permits straightforward weight concatenation and fusion. However, when considering full block parallelization entire Transformer blocks, each comprising both an attention module and an FFN, are run in parallel-such fusion is not possible. This is because each full block consists of two distinct components, with the attention output explicitly directed to its paired FFN. Running entire blocks completely in parallel not currently natively supported by heavily optimized inference frameworks such as TenSorRT-LLM or v vLLM (24). Nevertheless, one can envision assigning each full block to a different GPU, thereby maximizing parallelism and potentially achieving significant speedups. We use our block-wise dependency analysis to examine the dependency structure between blocks and identify candidate groups whose outputs are relatively independent. Such groups are ideal for parallelization with minimal accuracy loss and could enable higher degrees of tensor parallelism each full block operating on own set of GPUs-thereby improving inference throughput. our experiments, we implement these full block parallelization strategies using more flexible environment (e.g.. HuggingFace Transformers (46)., albeit with reduced optimization for large-scale deployments. thorough investigation of the actual speedups in specialized frameworks is left for future work.
15

---

(restricted to the set Two SwiGLU layers:
Single Fused SwiGLU:
Wi
Wi
W
W
W
da
dn
2d
W
W1 17
W2 13
W
2d4
4,
W3
W
24,-
W33
2T
FFN*
FFN + FFN = (X)
WET
WS
FFN (X)
W7"
N2*
Figure 7: A visualization of FFN Fusion applied to SwiGLU. Two FFNs (left) are fused into single FFN (right).
Method. We first compute the block-wise dependency matrix exactly as in Section 3 search for candidate sequences is mostly related to hardware constraints. Arbitrary sizes will be considered in our future work. For each such sequence [i, i + 3], i E [m - 4], we extract the corresponding M [8,4+3] E 18 1g44 submatrix from the block dependency matrix and compute two statistics:
We then consecutive blocks to fuse. The choice this number of blocks
M [+++3] = max M [,,+3] max
M 11.+3 sum
Σ KJEN
M
b.+43]
lower Mmax. and Msum indicates weaker dependencies among those blocks, suggesting they may be more amenable to parallelization. We exclude any blocks that lack attention (shaded regions in Figure (8b) because our focus here is on full Transformer blocks that contain both attention and FFN modules. We employ simple greedy algorithm to choose the best subsequences:
1. Choose the length-4 sequence starting index i" = arg min, M [6,+3]
2. If the argmin unique, break the tie using i* arg min, M [1++3) from step 1).
3. Parallelize the blocks according to the chose sequence [i", + +3).
4. Remove any overlapping sequences.
5. Repeat until no valid sequences remain.
The motivation behind this algorithm is to avoid choosing sequences with high dependency pairs. The sum tie breaking is due to step-like characteristic of Mmxx statistic (see Figure g7). Table 7 reports the downstream performance (MMLU and MT-Bench) after fusing the sequences chosen by 4 algorithm steps.
16

---

Results. From Table 7 we observe that par- allelizing the first and second sequence of four blocks leads to only modest drop in MMLUU, but once the third sequence is added, the de- cline becomes much more pronounced. This suggests that full block parallelization more challenging than FFN Fusion. These observa- tions are supported by the block-wise depen- dency heatmap and the submatrix statistics (Fig-
Table 7: Comparison of MMLU and MT-Bench
scores across different parallelization strategies.
Model Name
MMLU MT-Bench
49B model (no fusion)
80.73
8.87
Parallel Sequences
[38, 41
80.51
8.67
[38, 41
771.74
80.04
8.67
ure [). both of which indicate stronger inter- block dependencies and higher Mmax and Maum values in full block sequences compared to at- tention removed sequences. In addition, Figure significantly the tokens directions.
[38, 41
[71, 74
[32, 35
77.86
8.20
[38, 41
[71, 74]
[32, 35|
[26, 29]
56.12
6.62
11 shows that the full transformer blocks alter
Measured Block
10
Mesic
1.5
30
Drapped Block
30
40
50
0.5
60
70
80
Cosine Dist
(a)
(b)
Figure 8: (a) Block-wise Dependency Heatmap of the 49B model (log-scale). Darker blue hues indicate weaker dependencies, darker red hues indicate stronger dependencies. (b) Mm and M values for 4-Block Sequences of the 49B model. Lower values indicating more promising candidates for parallelization.
C Datasets
For applying Puzzle throughout our experiments, we used the same dataset mixture used in + termed Distillation Mix. This mixture includes source code repositories, Wikipedia articles, books, news websites, and several other domains. The dataset comprises 224 billion tokens collected from three public datasets: Web (34), Dolma [40], and Buzz-V1.2 [21] For UJtra-253B-Bar KD training we used the same data reinforced with synthetic data generated with Llama-405B, following the 48 approach, to help further align Mitra-233--Bass with its parent model.
D Ultra-253--Bass and 49B Architecture Overview
In this section, we detail the configuration of each block in the Ultra-253B-BB 353B-Base model and highlight several unique design choices introduced by the Puzzle framework on top of Llama-405B. Notably, the model employs variable FFN widths, with scaling multipliers ranging from 0.4875 up to 4.875 across its 126 blocks. This variability allows the model to dynamically adjust capacity at different depths, balancing performance with computational efficiency. Furthermore, it includes a sequence of 50 consecutive layers that bypass the attention mechanism (denoted by no_op), dedicating these layers entirely to FFN processing and creating ideal region for FFN Fusion.
We observe a similar pruning pattern the 49B model, which is derived from the Llama-3.1 70B model using the Puzzle framework. Like its larger counterpart, the 49B model also features variable FFN multipliers and blocks where attention removed, although at a smaller scale. Figure compares both architectures.
17

---

3.5
3
2.5
3
2
1.5
2
0.5
(a) UJtra-253B-Baa Attention Layout
(h) Ultra-253B-Bar Base FFN Layout
8
Sayar Type
3
3
2
20
40
60
80
0 20
40
60
80
(c) 49B Attention Layout
(d) 49B FFN Layout
Figure A 2 2 overview of Ultra-2533 Base (top row) and 49B model (bottom row). Subfigures and (b) illustrate the attention and FFN configurations, respettvely... for the 253B model. Subfigures and ( show the corresponding layouts for the 49B model. Both architectures feature variable FFN widths and regions where attention has been removed, although at different scales.
E MoE Inference Performance
MoEs show great promise in reducing inference costs and currently spearheaded by DeepSeek- V3 [6]. a highly sparse MoE with a great performance. However, MoEs have their own problems in inference.
E.1 Problems with Small Blocks
An underlying problem for MoEs paradoxically arises from the small size of their sub-modules. Smaller layers incur large latency than expected for two reasons:
GPU utilization low level GPU overheads such as wave quantization produce latency problems that become stronger as the module becomes smaller.
Communication overhead The Ael-Reduce operation that is typical for tensor parallelism creates an overhead for communication. For small modules this overhead increases in proportion.
MoE's experts are smaller than dense MLPs and hence suffer more strongly from the above effects. In addition the routing mechanism MoEs another small module that suffers accordingly. As result, dense models scale better with the number of GPUs while Mořs reach an early saturation. worth to notice that this in complete contrast to our FFN Fusing method that creates larger modules that parallelize better, visualized in Figure 2
18

---

E.2
Batch Size Effects
MoEs behave very well for very large batch sizes. For such large workloads load balancing issues are less dominant and low level overheads become negligible. However, smaller batch sizes are not as efficient. Concretely, we define an efficient size of a batch through the perspective of analyzing linear layer. The performance of linear layers varies significantly with batch size. For small batches, performance is limited by the time needed to read weights from DRAM (I/O bottleneck). For large batches, performance becomes compute-boune Notably, the latency of linear layer remains the same for batch sizes 1 and 64, as both are dominated by weight loading time. This makes small batches extremely inefficient. While other factors exist, we consider a linear layer to operate efficiently once the number of tokens the batch approaches 64, and inefficient otherwise. MoEs require all experts to be fully utilized for high throughput. in practice, the total batch is split across experts, significantly reducing the effective number of tokens each expert sees. For example, with a sparsity factor of 1/32 (as in [5] ()), batch of 2048 tokens (64 x 32) is needed to ensure each expert processes around 64 tokens. The case is more extreme in Llama 4 Maverick, where 1/128 sparsity factor requires batch of 8192 tokens (64 128) reach this threshold. These batch sizes are often impractical-espeeiilll in generation mode mode-leaving MoEs to operate in low-efficiency regime, where linear layers under-utilized. We can thus conclude that MoEs only really live up to their promise for very large batch sizes. For more commonly used intermediate batch sizes they suffer from bad scaling with the number of tokens, larger low level overheads, and worse parallelization scaling than dense models.
limited by the hardware's 's maximal FL.OPs/sec.
F Further Pairwise Block Dependency Analysis
Measured Block
20
60
80
100
120
40
Dropped Block
100
120
Cosine Dist
2
2
2
2
2 2
Figure 10: Replication of the block-wise dependency heatmap for Ultra-253B-Bas PreFusion, shown here for convenience. Each coordinate (i, represents the cosine distance between h (X) and Ti (X), quantifying how much block j depends on block i. Darker blue indicates weaker dependencies, while darker red indicates stronger dependencies. The dotted box marks attention- removed region with generally low dependency values, suggesting high parallelization potential.
Figure offerss a more detailed view of the block-wise dependency structure in JJtra-253B-- BBBe PreFusion. Below, we highlight several notable observations:
Overall Color Scale. Each entry (i,j) in the heatmap corresponds to the cosine distance between the outputs of block j when block is dropped versus when the model is intact. A dark red hue signifies a strong dependency, meaning that omitting block substantially alters block j. Conversely, dark blue indicates a weak dependency, suggesting that dropping block does not notably affect block j.
19

---

Aetention-emooved Region (Dotted Box). The dotted box highlights region of consis- tently low dependency values, where most blocks exhibit minimal mutual influence. This area arises from attention pruning in those layers and is particularly well-suited for FFN Fusion, since layers here can more easily parallelized without significantly harming accuracy.
Crucial Early Blocks. Some blocks in the upper rows of the matrix (toward the top of the y-axis) display consistently red cells across many columns. This indicates that certain early layers have a strong influence on a wide range of subsequent blocks, making them less amenable parallelization.
Sub-region with Sequential Reliance. A noticeable diagonal band of higher values appears in some sub-regions, implying that each block that band strongly depends on its immediate predecessor. Such sequential reliance reduces the potential for parallelization in those areas, since omitting any single block significantly affects the next one.
Global Sinks at Later Layers. As we move toward deeper layers (near the bottom-right corner, some blocks again show stronger dependencies, acting "global sinks" that consolidate information from many earlier blocks. Although the importance of certain blocks may wane in mid-network regions, it can resurface in the final layers.
In summary, the dependency matrix reveals both strongly and weakly coupled sub-regions. High- dependency zones require careful consideration to avoid accuracy loss if parallclization is attempted. Conversely, low-dependency zones, such as the attention-remove region, offer a promising avenue for efficient FFN Fusion.
G Fusion Explainability
20

---

0.4
0.3
0.25
0.1
0.05
0.6
0.5
0.3
0.2
15
20 30
70
30
90 100
(C)
Figure 11: Per-layer metrics. Upper row is the cosine distance between f(X) and X for the (a) The 49B model and (b) U5tB--253B-- BBsse model. Bottom row represents the ratio between h(X) and X for the (c) The 49B model and (d) Ultra-553B--as model.
NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The claim introduced in the abstract and introduction are supported by theoret- ical claims (Theorem 3.1 and empirical evidance (Sections 4 & 5)
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
21

---

Answer: [Yes]
Justification: The authors indicate that the suggested method is mainly relevant in cases where there is long a sequence consecutive FFN blocks.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: Theorem 3.1 proves the equivalence between parallel computation of FFN
blocks and our consolidate formulation.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
22

---

Justification:
The steps that were taken to create the results are fully described in our method & results sections.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which authors are welcome to deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [NA]
Justification:
We will release code upon acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.ce)
wtw.aubmiissiooooooooooooooo icy for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https //nips cc/
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
23

---

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: The paper describes all the details regarding the main method (FFN Fusion). Training methods were based on previous work.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [NA]
Justification: Error bars are irrelevant for the experiments in the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
24

---

Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental rums well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (c.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
25

---

Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the paper does not use existing assets. The authors should cite the original paper that produced the code package dataset.
The authors should state which version the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/codeemoodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
26

---

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy for what should or should be described.
wwwpurripssc..ooofeeeeeeeeeeeeee /
27


Reviews:
[{"review_id": "eYkb2MxkJz", "reviewer": "Reviewer_RXRj", "summary": "This paper presents FFN fusion, a post-training architecture optimization technique that removes attention layers from certain Transformer blocks and fuses the remaining feedforward layers into one. The method is novel and sheds light on the effectiveness of individual attention/feedforward layers in a pretrained LLaMA model. The downstream performance along with inference speedup appears promising.", "questions": "To what extent do you think the current method can generalize? One concern I have is the normalization strategy. Specifically, the usage of pre-norm might affect many design choices in this paper. For example, the h(x) function used in the pairwise block dependency seems to depend on the pre-norm structure. In cases with other normalization structures like post-norm, can we still remove attention layers and merge FFN layers effectively?", "limitations": "Strong performance with unclear generalization across different architectures and settings.", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "### Strength\n1. Well-motivated and well-designed approach that provides valuable insights into layer-wise contributions in pre-trained models.\n2. Strong performance demonstrating improvements in both quality and efficiency metrics.\n\n### Weaknesses:\nMy main concern lies in the generalization of this method:\n\n1. It's unclear how the method can generalize to MoE models, which are standard architectures for very large models.\n2. The coupling between the method and LLaMA-specific design choices remains unclear (see question below).\n3. The impact of removing attention layers on long-context performance is not addressed, though attention layers are presumably more crucial in long-context settings.", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive feedback,\n\n*”The impact of removing attention layers on long-context performance is not addressed, though attention layers are presumably more crucial in long-context settings.”*\n\nThis is an important point, and it is partially addressed in Figure 4, showing that the RULER [1] long-context benchmark results at 128K shows that our 253B CPT model, with half the attention layers of the original Llama-405B, achieves an average score across all RULER tests that is 112.85% of Llama-405B’s. In addition, even before any training or knowledge distillation (KD), the FFN-fused model retains 98.7% of Llama-405B’s performance at 64K context and 79.4% at 128K. We will highlight these results in the revised manuscript, to clarify that long-context performance is preserved.\n\n*”One concern I have is the normalization strategy. Specifically, the usage of pre-norm might affect many design choices in this paper. For example, the h(x) function used in the pairwise block dependency seems to depend on the pre-norm structure. In cases with other normalization structures like post-norm, can we still remove attention layers and merge FFN layers effectively?”*\n\n\nMost current LLMs adopt a pre-norm architecture, including models like Llama, Mistral, DeepSeek, Qwen, and Kimi among many others. Our work focuses on this prevalent setting. That said, exploring the applicability of our approach to post-norm architectures is an interesting direction for future work.\n\nWhile we were unable to conduct a full empirical study during the rebuttal period, we can offer a structural analysis that highlights potential differences. The main distinction is that post-norm applies normalization after the block output $f(X)$, in contrast to the accumulating behavior seen in pre-norm models. In pre-norm architectures, we observe that token representation norms tend to increase significantly across layers, which results in a decreasing cosine distance between $X$ and $f(X)$—except in the final layers. These lower cosine distances appear to correlate with successful attention pruning and FFN fusion.\nIn post-norm models, we expect a different pattern of cosine distances due to the normalization applied at the output. This could lead to different behavior in terms of attention removal and fusion. Nonetheless, we believe that the ability to reduce model capacity—either by pruning attention or fusing FFNs—stems primarily from the over-parameterization of LLMs, rather than from specific architectural choices. For instance, [3] show that early layers in post-norm models are typically less effective, which indicates they might also tolerate attention pruning and FFN Fusion.\n\n\n*”It's unclear how the method can generalize to MoE models, which are standard architectures for very large models.”*\n\nThis is indeed a very interesting future direction! We have already developed such possible generalizations of the FFN fusion method for MoE architectures. The simplest of these, involving two MoE layers with T total experts, s shared experts, and r routed experts from which we select k chosen experts per layer, proceeds as follows: first, attention is removed between the MoE layers (where applicable); then, the routers are concatenated, producing an output of dimension 2r. A top-K selection then chooses 2k active experts from the list of 2r experts, and the shared experts are fused into a 2x larger shared expert. \nOur preliminary evidence suggests that some pairs of Llama 4 Scout MoE layers can be fused in such a manner without loss of accuracy in MMLU and MT-Bench. We will explore the effectiveness and refinement of these ideas in future work, and will add this point to the concluding remarks section.\n\n\n[1] RULER: What's the Real Context Size of Your Long-Context Language Models? Hsieh et al.\n\n[2] Puzzle: Distillation-Based NAS for Inference-Optimized LLMs, Bercovich et al.\n\n[3] The Curse of Depth in Large Language Models, Sun et al."}, {"review_id": "h33dO7K19I", "reviewer": "Reviewer_Pynj", "summary": "This paper identifies inference inefficiencies of recent Mixture of Expert architectures for small and medium batch sizes due to the multiple small FFN (sub-)modules (i.e. experts).\nInstead, the paper improves hardware utilization during inference by combining FFN layers into single wider layers after training, thereby reducing the serial computation (due to the number of layers) and increasing the parallel computation. \nUsing the proposed method FFN Fusion, the paper introduces a 253B parameter dense model derived from Llama-3.1 405B, which shows a 1.71x inference speed up at equal or higher quality compared to the base model.", "questions": "- L. 168: For TP can‘t we overlap computation and communication and in this way hide the communication time?\n- Section E: The authors only talk about „single token inference“, so I assume the authors mean the inefficiency during generation (?), but it is not mentioned. How is it for prefill?\n- Section E: What are intermediate / large batch sizes in this case? The motivation would be much stronger with (a) concrete example(s).\n- Section E: What is meant with suboptimal and optimal region of efficiency? The explanation is rather vague.\n\n- Theorem 3.1 shows that the parallel execution is identical to the parallel execution of FFN layers. However, does it also hold in the presence of skip connections and norm layers around the sequential FFN layers?\n\n- Figure 2: Why is there a significant difference in the columns 110 - 120 compared to the earlier layers?\n- How do you compute the cosine distance over matrices n x d ? Do you average over the sequence dimension?\n\nSection 5.3: Do you have an explanation why the last layer in a sequence of blocks is especially important?", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- FFN Fusion yields higher inference throughput at equal or better performance than the base model.\n- Thorough empirical investigation (even though only at very large model sizes (70B and 253B)).\n\nWeaknesses:\n- The motivation of inefficient mixture of experts for small batch sizes would be stronger if the authors could give a concrete example including latency and throughput calculation or measurements, which demonstrate the weaknesses of Mixture of Experts (See also the questions on Section E).\n- The experiments are performed only on two models from the same family, which likely used the same training recipe and training dataset. While the authors provide thorough empirical analyses on these two models, it still leaves the question of how general this method is. Also, FFN fusion has only been applied to very large models.\n- Experiments at smaller scales (e.g. 8B - 32B), where there are more open source models available, could further prove the generality of FFN Fusion.\n- Figure 5: To really explore the efficiency frontier more models should be added to this plot (potentially including MoE models, to really see the difference in Latency, see also first point)", "was_revised": true, "final_justification": "The authors provided new insightful experiments, which address my concerns on generality (e.g. smaller models and/or models from other family) and they also clarified their motivation (provided in Appendix E).", "rebuttal": "Thank you for your feedback and helpful questions regarding Appendix E.\n\n*”L. 168: For TP can‘t we overlap computation and communication and in this way hide the communication time?”*\n\nTP does not enable communication hiding. The all-reduce (summation) operation at the end of the TP cycle requires all GPUs to complete their computation before communication can begin.\nIn theory, overlapping computation and communication is possible by partitioning data (activations) in some way. However, this is only practical through micro-batching, which introduces its own challenges. A fine-grained approach—such as streaming portions of the output as they become ready—shows little to no benefit when the data volume is small.\n\n*”Section E: The authors only talk about „single token inference“, so I assume the authors mean the inefficiency during generation (?), but it is not mentioned. How is it for prefill?”*\n\nWhile the MoE analysis focuses on generation, the key factor for inference performance here is not batch size but the number of tokens in the forward pass. In that sense, prefill behaves similarly to generation when the number of tokens is the same. For example, a prefill with 256 tokens has similar behavior to generation with a batch size of 256. So the analysis in Appendix E applies in both cases.\n\n*”Section E: What are intermediate / large batch sizes in this case? The motivation would be much stronger with (a) concrete example(s).”*\n\n*”Section E: What is meant with suboptimal and optimal region of efficiency? The explanation is rather vague.”*\n\nYour questions helped us clarify Appendix E, and we will include the following explanation in the revision:\n\n**Background: Linear layers.**\n The performance of linear layers varies significantly with batch size. For small batches, performance is limited by the time needed to read weights from DRAM (I/O bottleneck). For large batches, performance becomes compute-bound, limited by the hardware's maximal FLOPs/sec. Notably, the latency of a linear layer remains the same for batch sizes 1 and 64, as both are dominated by weight loading time. This makes small batches extremely inefficient. While other factors exist, we consider a linear layer to operate efficiently once the number of tokens in the batch approaches 64, and inefficient otherwise.\n\n**Back to MoEs:**\nAs mentioned in Appendix E, MoEs require all experts to be active for high throughput. However, the total batch is split across experts, significantly reducing the effective number of tokens each expert sees. For example, with a sparsity factor of 1/32 (as in DeepSeek-V3\\R1), a batch of 2048 tokens (64x32) is needed to ensure each expert processes around 64 tokens. The case is more extreme in Llama 4 Maverick, where a 1/128 sparsity factor requires a batch of 8192 tokens (64x128) to reach this threshold. These batch sizes are often impractical—especially in generation mode—leaving MoEs to operate in a low-efficiency regime, where linear layers are under-utilized.\nWe loosely define the *efficient* or *optimal* zone as having ≥64 tokens per expert; below that, operation enters a *suboptimal* or *inefficient* region due to underutilized linear layers.\n\n*”The experiments are performed only on two models from the same family”*\n\nPoint well taken. During the rebuttal period, we added experiments on Llama 8B 3.1 (see below) from the same family, as well as applying FFN fusion to Mistral Large 2, expanding beyond the Llama family. The process for the Mistral experiment was as follows:\n\n(1) **Applying the Puzzle algorithm** to remove attention layers. We removed 27 layers that resulted in 2 sequences of 14 and 13 consecutive FFN layers.\n\n(2) **Applying FFN Fusion** to fuse all but the final FFN layer in each sequence.\n\nWe did not have time during the rebuttal to run any recovery training, but for the revision we will also include results with KD.\nLooking at the final results, the fused model retained 99.4% of the pre-fusion MMLU score and 94.4% of its MT-Bench score.\nWe also intend to include models from additional families for the revision.\n\n\n\n*”Experiments at smaller scales (e.g. 8B - 32B), … could further prove the generality of FFN Fusion.”*\n\n\nTo determine if FFN Fusion is effective for smaller models, we tested it on the Llama 8B 3.1 Instruct model. The process involved:\n\n(1) **Applying the Puzzle algorithm** to remove attention layers. We removed 9 layers, resulting in 8 consecutive FFN layers—a slightly smaller proportion compared to the 70B case.\n\n(2) **Applying FFN Fusion** to fuse all but the final FFN layer in this sequence.\n\n(3) **Running a short KD recovery phase** over 20B tokens.\n\nCompared to the original (unmodified) Llama 8B model, the final fused model retained 98.99% / 96.5% of its performance on MMLU / MT-Bench, respectively. This confirms that FFN Fusion is effective at smaller scales as well. \n\n*”Figure 5: To really explore the efficiency frontier more models should be added to this plot (potentially including MoE models, to really see the difference in Latency, see also first point)”*\n\n We agree, and plan to add models such as Mistral and Deepseek in the revision.\n\n\n*”Theorem 3.1 shows that the parallel execution is identical to the parallel execution of FFN layers. However, does it also hold in the presence of skip connections and norm layers around the sequential FFN layers?”*\n\nTheorem 3.1 shows that Equation (5), which describes a parallel execution of FFN layers with shared normalization parameters, is mathematically equivalent to a single denser FFN layer using the same normalization parameters. Our main empirical observation is that in trained LLMs, sequences of FFN layers without intermediate attention—though each with their own skip connections and normalization parameters—can still be replaced with the parallel form in Equation (5). As you noted, while this substitution from a sequential to a parallel architecture isn't mathematically equivalent, our results prove it to maintain accuracy while improving efficiency.\n\n*”Figure 2: Why is there a significant difference in the columns 110 - 120 compared to the earlier layers?”*\n\n\nWe observed a distinct change in dependency patterns between the fused FFN layers (layers 66–114) and the subsequent layers (114–120). Unlike the earlier sequence, the later layers retain their attention modules (they were not removed by Puzzle) and show stronger dependencies along the main diagonal.\n\nWhile we don’t have a definitive explanation for this phenomenon ([1] attempts to characterize different functional stages of the network), this increased dependency seems to make the later layers harder to modify. Although we lack direct evidence, we hypothesize that these final layers play a key role in shaping the token representations for the LM head and final token prediction.\n\n*”How do you compute the cosine distance over matrices n x d ? Do you average over the sequence dimension?”*\n\nYes, we compute the cosine distance across the hidden dimension for each token, and then average the results over the sequence dimension.\n\n\n*”Section 5.3: Do you have an explanation why the last layer in a sequence of blocks is especially important?”*\n\nThat’s an interesting question. While we don’t have a definitive answer, we hypothesize that the final FFN in a sequence may serve a special role. In our case, Puzzle chose not to remove the attention layer following it—possibly indicating that this attention depends on the output of the preceding (non-fused) FFN.\n\n[1] The Remarkable Robustness of LLMs: Stages of Inference?, Lad et al."}, {"review_id": "AJBuxHjfkt", "reviewer": "Reviewer_zYpm", "summary": "This paper proposes FFN Fusion, a novel optimization method that removes attention layers from transformer-based LLMs and fuses sequential FFNs so that they are processed in parallel, leading to significant memory and latency improvements. The method operates by leveraging cosine distance to identify the preceding blocks of which a given block depends on. It then removes the attention blocks of contiguous blocks with low inter-dependency and fuses the remaining sequential FFN layers through a theorem that leverages the use of SwiGLU in modern LLMs. Empirical results show that this method has a relatively minimal performance impact as compared to the original model, and additional post-application finetuning leads to even stronger performance. Ablations show that removing the FFN layers entirely  can severely impact performance and that reversing the FFN layers has little impact, suggesting that the small magnitudal and directional changes of each FFN are not inter-dependent.", "questions": "- Is the proposed method effective for models outside of the Llama family?\n- what was the time/compute usage for the experiments?\n- how does the proposed method par against existing quantization/pruning methods?\n- Do we still see similar patterns of inter-block dependency with smaller models?\n\nI would be happy to raise my score to Accept if these questions are effectively and thoroughly addressed.", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Strengths:\n\n- The initial empirical results strongly support the effectiveness of the method\n- The method is applicable to modern LLMs, for which there is a significant need for techniques of this nature\n- the work conducts ablations which both demonstrate that the FFNs are still needed and provide explanation for why the method works\n- the paper is generally well written with a clear motivation for why the method is useful\n\nWeaknesses:\n\n- it appears that only two models were evaluated, and they both came from the Llama family.  It is therefore unclear whether this method will generalize to models outside of the Llama family that use SwiGLU.\n- There is no discussion on how much time/compute was used for the pairwise block dependency calculations\n- While the motivation discusses alternate methods like quantization and pruning, there is no discernable comparison to techniques from these fields. It would be nice to see these comparisons. Perhaps the best technique from each field which can equivalently bring down the memory/latency as FFN Fusion does, and how the performance compares.\n- While 70B and 405B are the most relevant model sizes for this technique, it would have been interesting to see even smaller sizes and whether or not the insights found for the larger sizes generalize to smaller ones. Because there is no discussion of time/compute needed for this method, I am not sure if this is a reasonable request for the discussion period", "was_revised": true, "final_justification": "The authors sufficiently addressed all of my concerns, of which there were many, with strong empirical results and analysis. The proposed method is novel and strongly backed by empirical results. It adds a new compression technique that is both complementary and combinable with existing techniques in the field, making it both relevant and practical.", "rebuttal": "Thank you for your detailed feedback,\n\n*”it appears that only two models were evaluated, and they both came from the Llama family.”*\n\nPoint well taken. During the rebuttal period, we added experiments on Llama 8B 3.1 (see below) from the same family, as well as applying FFN fusion to Mistral Large 2, expanding beyond the Llama family. The process for the Mistral experiment was as follows:\n\n(1) **Applying the Puzzle algorithm** to remove attention layers. We removed 27 layers that resulted in 2 sequences of 14 and 13 consecutive FFN layers.\n\n(2) **Applying FFN Fusion** to fuse all but the final FFN layer in each sequence.\n\nWe did not have time during the rebuttal to run any recovery training, but for the revision we will also include results with KD.\nLooking at the final results, the fused model retained 99.4% of the pre-fusion MMLU score and 94.4% of its MT-Bench score.\nWe also intend to include models from additional diverse families for the revision.\n\n\n\n*”it would have been interesting to see even smaller sizes and whether or not the insights found for the larger sizes generalize to smaller ones”*\n\nTo determine if FFN Fusion is effective for smaller models, we tested it on the Llama 8B 3.1 Instruct model. The process involved:\n\n(1) **Applying the Puzzle algorithm** to remove attention layers. We removed 9 layers, resulting in 8 consecutive FFN layers—a slightly smaller proportion compared to the 70B case.\n\n(2) **Applying FFN Fusion** to fuse all but the final FFN layer in this sequence.\n\n(3) **Running a short KD recovery phase** over 20B tokens.\n\nCompared to the original (unmodified) Llama 8B model, the final fused model retained 98.99% / 96.5% of its performance on MMLU / MT-Bench, respectively. This confirms that FFN Fusion is effective at smaller scales as well. \n\n*”Do we still see similar patterns of inter-block dependency with smaller models?”*\n\nThat’s a good question. Beyond the fusion experiment for Llama 8B 3.1 Instruct, following your question, we also analyzed the model's inter-block dependency. We observed that, like the 253B and 49B models, the fused layers showed low dependency, while other layers were more dependent on their neighbors. We aim to add the plot to the revised version alongside the corresponding fusion results.\n\n*”There is no discussion on how much time/compute was used for the pairwise block dependency calculations”*\n\n*”what was the time/compute usage for the experiments?”*\n\nCalculating the pairwise block dependencies took 0.24 H100 GPU hours per matrix row for Llama 405B.\nIt’s important to clarify that this step was part of an exploratory experiment and is not required for applying FFN Fusion itself. The fusion step is extremely lightweight—it involves direct merging of FFN matrices and does not require any heavy computation. This is why we did not report compute usage for it. The block dependency analysis is only relevant to the Block Parallelization method discussed in Appendix B.\nOther compute-intensive steps, such as KD training or CPT, are detailed in Section 4 under “Additional Training”, where we provide the full token budgets.\nAre there any other time/compute steps you think we should also include in the revision?\n\n\n*”While the motivation discusses alternate methods like quantization and pruning, there is no discernable comparison to techniques from these fields. It would be nice to see these comparisons.”*\n\nOur paper presents FFN Fusion as a technique complementary to existing inference efficiency methods like pruning and quantization. Pruning improves matrix multiplication speed by leveraging sparsity, while FFN Fusion reduces synchronization overhead by decreasing the number of computational blocks.\nNotably, the models in this paper already combine both approaches: FFN channel pruning (applied during Puzzle) and FFN Fusion, showcasing the benefits of their combined application. Similarly, all results in Section 4 are shown with FP8 quantization. Importantly, our experiments (detailed in Section 5.2) already compare FFN Fusion to pruning entire FFN layers, which we feel is the best baseline for our method, and have demonstrated the superior performance of our approach.\n\nHowever, to offer another comparison between FFN Fusion and pruning, we ran an experiment on the 49B model from Section 5. We applied structured sparsity (Wanda [1]) to the FFN weights, achieving 40% sparsity. The resulting model had a latency profile similar to “step 3” in Table 3. The sparse model scored 78.26 MMLU, while our fused model achieved 80.36 MMLU. As a next step—also planned for inclusion in the revision—we aim to explore combining both approaches.\n\n[1] A Simple and Effective Pruning Approach for Large Language Models, Sun et al.\n\n\nWe hope we have thoroughly addressed your questions."}, {"review_id": "YpZ8qxvghb", "reviewer": "Reviewer_Cc25", "summary": "The paper proposes a modification to the existing LLM architecture by: 1) removing some attention layers (using an existing method called Puzzle), and 2) fusing sequential FFNs to improve the runtime efficiency of LLMs. The authors apply these techniques to Llama-3.1-405B-Instruct and create a 253B model, referred to as 253B-Fusion, which achieves a 1.71× inference speedup while maintaining accuracy. Knowledge distillation is used to recover model quality.", "questions": "1. The majority of the performance gains come from the existing Puzzle method rather than Fusion.\nTo construct the fused model, the authors first derive a pre-fusion 253B model by applying Puzzle to Llama-405B. From Table 2, this model already achieves a 1.53× speedup over the original 405B model, primarily due to reduced model size and the elimination of attention layers. The current work proposes further speedup by fusing sequential FFNs into a wider FFN to enhance GPU utilization. However, this yields only about a 12% additional speedup, which is relatively minor considering it is a lossy transformation.\n\n2. The speedup from fusing large matrix multiplications offers limited potential.\nIn large models, the FFN layers themselves are quite large. As GPU memory capacity increases, larger weight slices can be stored on each GPU under tensor parallelism. This reduces the benefits of fusion, thereby limiting its potential for further speedup.", "limitations": "No", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n1. The topic of exploring runtime-efficient LLM architecture is important.\n2. The authors provide solid experimental results validating the fused model’s accuracy.\n\n**Weaknesses**\n1. The majority of the performance gains come from the existing Puzzle method rather than the proposed Fusion technique.\n2. The speedup from fusing large matrix multiplications offers limited potential.", "was_revised": true, "final_justification": "The response of the authors address some of my concerns (the second weakness) but does not address the first weakness in my review. But I am in general positive for this paper. I tend to keep my weak accept score.", "rebuttal": "Thank you for your positive feedback,\n\n*”The speedup from fusing large matrix multiplications offers limited potential. In large models, the FFN layers themselves are quite large. As GPU memory capacity increases, larger weight slices can be stored on each GPU under tensor parallelism. This reduces the benefits of fusion, thereby limiting its potential for further speedup.”*\n\nFFN Fusion delivers *larger* speed‑ups on newer and better GPUs. To empirically demonstrate this, we compared the original Llama‑405B and our 253B‑Fusion on TP 8 B200 nodes, observing a 1.83× latency improvement, which exceeds the 1.71× measured on TP 8 H100.\n\nIf we understand your argument correctly, you suggest that increased GPU memory in newer GPUs would allow running with lower TP (otherwise, with the same TP the same-sized weight slices would be stored on each GPU), but this is not necessarily true.\nA complete cost model must also weigh communication. **Latency**‑sensitive inference typically runs at small batch sizes where high‑bandwidth memory (HBM), not total memory DRAM, is the bottleneck. Even when the entire model fits on one newer GPU—for example, 253B‑Fusion at FP4 on a single B100 or B200—running in TP to shorten the critical path of reading weights from HBM and to share the workload across multiple devices is still preferable. At these small batch sizes the I/O of fetching weights, not raw compute, limits latency, so only greater HBM bandwidth or distributing the read over several GPUs helps; simply adding more DRAM capacity does not. Using TP, however, inserts a synchronization for every layer, which can be eliminated using FFN Fusion. Thus the relative advantage of FFN Fusion increases as TP deepens and as precision is reduced, because both trends amplify communication costs.\n\nTo illustrate this point, we benchmarked 253B-Fusion with FP4 weights on a B200 GPU, and compare it to the 253B-Fusion measurements for FP8 in the manuscript. With no TP, a single B200 run is 0.53× slower than TP 8 on H100, highlighting the HBM’s bandwidth ceiling. Moving to TP 2 on B200 closes the gap to 0.98×, and at TP 8 the same model runs 2.8× faster than TP 8 on H100.\n\nFurthermore, FFN Fusion is designed to improve practical aspects of performance that arise from inefficient matrix multiplications (GEMMs) for matrices that are too small compared to the GPU capabilities (See Appendix E.1.), and to alleviate inter-device communication overheads by reducing the number of layers. \nEach FFN matrix looks “large” on paper, yet under TP it is sliced into tiles that may occupy only a fraction of a GPU’s streaming multiprocessors (SMs). Let us consider a practical example. In Llama 70B each MLP weight matrix is 8K*28K or 28K*8K (224MB in total assuming FP8 data type). Running with TP8 each GPU weight matrix is sharded to just 8K*3.5K or 3.5K*8K (224MB/8=28MB). The sizes (and number of tiles) discussed are too small to saturate even an H100. Because B200 devices have even more SMs, and because systems such as NVL72 allow TP factors well beyond eight (up to TP72), the under‑utilization would become even worse. For these reasons, matrices that appear \"large\" are often still too small to utilize hardware efficiently—making FFN fusion appealing, especially on future GPUs.\n\nIf we misunderstood your question or if our explanation didn’t fully address it, we’d be happy to clarify further.\n\n*”The majority of the performance gains come from the existing Puzzle method rather than the proposed Fusion technique.”*\n\nWe agree that most performance gains for 253B-Fusion were achieved by Puzzle. Importantly, Puzzle and FFN Fusion are complementary techniques—their benefits stack rather than compete. Puzzle is a collection of efficiency methods (including channel pruning, layer removal, operation replacements like GQA, etc.), and FFN Fusion adds another orthogonal tool to the toolbox. Also, as noted earlier, the speedup from FFN Fusion can increase with stronger GPUs or larger TP configurations.\n\nLooking ahead and after recognizing FFN Fusion’s potential benefits, one could also consider architectural decisions and NAS strategies to further increase the potential benefits:\n\n(a) Longer FFN sequences could be promoted using NAS techniques like Puzzle. The current 253B model was produced using Puzzle’s default settings. However, Puzzle could be adapted to favor attention removals that later enable fusion—making FFN Fusion even more impactful. In fact, FFN Fusion could be introduced as an “alternative block” type within Puzzle, allowing the MIP solver to jointly consider attention removal and FFN fusion during optimization. We have begun preliminary experiments in this direction.\n\n(b) Most current LLMs are built with homogeneous transformer blocks, each consisting of an attention layer and an FFN layer. However, future models could be designed with sequential FFN layers from the start, increasing the number of FFN sequences that could potentially be fused. This could even benefit model performance, as demonstrated in [1].\n\n\n[1] Improving Transformer Models by Reordering their Sublayers, Ofir Press et al., ACL 2020."}]

Author rebuttal:

Meta-review: This paper proposes FFN fusion, an approach that parallelizes sequential layers in a Transformer by removing specific attention layers and running subsequent FFN layers in parallel. The approach reduces inference latency while preserving model behavior, and the authors evaluate it on Llama-3.1-405B.

The reviewers found the problem space important and found the method to be well-motivated and well-designed, with strong empirical results and baselines.

Given the importance of the problem and novelty in the approach, I recommend this paper for accept with oral or spotlight.
Decision: Accept (spotlight)