First Attentions Last: Better Exploiting First Attentions for Efficient Parallel Training
Authors: Gyudong Kim
Hyukju Na
Jin Hyeon Kim
Hyunsung Jang
Jaemin Park
Jaegi Hwang
NAMKOO HA
Seungryong Kim
Young Geun Kim
Venue: NeurIPS 2025 poster ()
Keywords: Neural Networks
Parallelism
Distributed Training
Tensor Parallel
Primary area: deep_learning
TL;DR: New Transformer architectures, FAL and FAL+, reuse the first attention output carefully —FAL cuts TP communication for up to 44% faster training, while FAL+ improves perplexity.

Abstract: As training billion-scale transformers becomes increasingly common, employing multiple distributed GPUs along with parallel training methods has become a standard practice. However, existing transformer designs suffer from significant communication overhead, especially in Tensor Parallelism (TP), where each block’s MHA–MLP connection requires an all-reduce communication. Through our investigation, we show that the MHA-MLP connections can be bypassed for efficiency, while the attention output of the first layer can serve as an alternative signal for the bypassed connection. Motivated by the observations, we propose FAL (First Attentions Last), an efficient transformer architecture that redirects the first MHA output to the MLP inputs of the following layers, eliminating the per-block MHA-MLP connections. This removes the all-reduce communication and enables parallel execution of MHA and MLP on a single GPU. We also introduce FAL+, which adds the normalized first attention output to the MHA outputs of the following layers to augment the MLP input for the model quality. Our evaluation shows that FAL reduces multi-GPU training time by up to 44%, improves single-GPU throughput by up to 1.18×, and achieves better perplexity compared to the baseline GPT. FAL+ achieves even lower perplexity without increasing the training time than the baseline. Codes are available at: https://casl-ku.github.io/FAL/

First Attentions Last: Better Exploiting First Attentions for Efficient Transformer Training
Gyudong Kim
Hyukju Na
Jin Hyeon Kim
Hyunsung Jang
Jaemin Park
Jaegi Hwang
Namkoo Ha
Seungryong Kim
Young Geun Kim
Korea University
LIG Nexl Co., Ltd.
$ KAIST AI
Abstract
As training billion-scale transformers becomes increasingly common, employing multiple distributed GPUs along with parallel training methods has become a standard practice. However, existing transformer designs suffer from significant communication overhead, especially in Tensor Parallelism (TP), where each block's MHA-MLP connection requires an all-reduce communication. Through our investigation, we show that the MHA-MLP connections can be bypassed for efficiency, while the attention output of the first layer can serve as an alternative signal for the bypassed connection. Motivated by the observations, we propose FAL (First Aftentions Last), an efficient transformer architecture that redirects the first MHA output to the MILP inputs of the following layers, eliminating the per-block MHA-MLPP connections. This removes the all-reduce communication and enables parallel execution of MHA and MLP on single GPU. We also introduce FAL+, which adds the normalized first attention output to the MHA outputs of the following layers to augment the MLP input for the model quality. Our evaluation shows that FAL reduces multi-GPU training time by up to 44%, improves single-GPU throughput by up to 1.18x, and achieves better perplexity compared to the baseline GPT. FAL+ achieves even lower perplexity without increasing the training time than the baseline. Codes are available at: https://eus-ku.ggttuu......
Introduction
As transformers continue to grow following scaling law trends [1, Large Language Models (LLMs) such as GPT [2] and LLaMA 3 demonstrate superior performance across wide range of NLP tasks. Given that the large transformers usually have billions of parameters which far exceeding a single GPU's memory and computation limit, distributed training multiple GPUs is necessary.
Among distributed training methods, Tensor Parallelism (TP) 4 has been considered as a standard practice for multi-GPU training in each single-node, given its decent computation throughput and memory efficiency S 7 However, the overall efficiency of TP can still be limited by the communication overhead across the GPUs. For example, if the number of GPUs increases (and/or the interconnect speed slows down), the efficiency of TP can be significantly degraded due to the increased communication overhead. Hence, to further scale up transformers with more GPUs and diverse interconnect configurations, it is essential to mitigate the communication overhead TP.
One of the major sources of communication overhead in TP is the per- block communication (i.c., all-reduce) required across GPUs transfer data between two main modules in each transformer block: Multi Head Attention (MHA) and Multi Layer Perceptron (MLP). In TP, each GPU computes a part of the MHA output, but the results must be aggregated via all-reduce to form complete activation before being passed to the MLP.
Corresponding authors (seungryong kim @ kr. younggeun kim korea.ac.kr)
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Block Indes
Block index N
Computation
Commumication Perplexity
Nornalized Latency
First Alluniion
17.8
inserr 22
Tinear @
(b) FAL
17.6
Concert
Neciicarity Linerr
Black index e
Block Index N
17.4
Scaled Dat-Penduct Atteniion
0.2
17.2
K
Q T
Lineer
Multi-Layer Perceptron
17
Multi-Hes Head Attention
Lawarr Norm 5
GPT2 FAL FAL+
(a) Pre-I. Transformer
(c) FAL+
(d) Performance
Figure 1: Transformer block designs highlighting our proposed modifications. (a) Pre-LN archi-
tecture with Layer Normalization (LN), Multi Head Attention (MHA) and Multi Layer Perceptron (MLP). (b) FAL blocks with reconfigured connections, (c) FAL+ blocks with augmented connections. (d) End-to-end training time and perplexity comparison of Pre-LN architecture, FAL, and FAL+.
To reduce communication overhead in TP, we question whether the direct connections between two main modules in each transformer block (i.e., MHA-MLP connections) are strictly necessary. To answer this question, we conduct series of analyses (Sec. 3), revealing two key observations:
1. MHA-MLP Connections Can Be Reconfigured: We find that the direct MHA-MLP connection not always essential, as the residual path already accumulates MHA outputs from prior blocks. However, naively removing the connection causes a large drop in model quality, indicating additional mechanisms are required mitigate information loss.
2. First Attention is Key: Our analysis shows that the first MHA output has a disproportionnaely large impact on final predictions. Thus, leveraging the first attention more effectively can compensate for the information loss caused by the removed connections.
Motivated by findings, we propose FAL (Firs! Attentions Last), efficient parallel transformer architecture that redirects the high-impact first MHA output to the MIP inputs of following blocks, rather than relying on the attention within each block (Fig. [ (b)). This modification eliminates expensive all-reduce communication within each block and also enables parallel execution of MHA and MLP on single GPU. As result, FAL improves training throughput by 1.07-1.52x in multi- GPU settings, and by 1.03-1.18x x on single GPU, compared to the standard transformer. Moreover, carefully leveraging the high-impact first attention not only preserves model quality but improves reducing validation perplexity and increasing the average SuperGLUE (8) score. To further improve the model quality leveraging the high-impact first attention, we additionally propose FAL+, a variant of FAL, which augments MHA-MLP connections rather than removing them (Fig. (c)) FAL+ achieves even lower perplexity than FAL. further confirming the importance of the first attention.
Background
Transformer Architecture. A transformer model consists of stack of transformer blocks. Fig. (a) depicts diagram of Pre-LayerNorm (Pre-LN) architecture, which is widely adopted for deep transformers due to its superior training stability and signal propagation [9] IO, [I]. Each transformer block includes Multi Head Attention (MHA) for dependency modeling and a Multi Layer
Perceptron (MLP) for position-wis wise transformation using the dependencies computed by MHA.
The MHA first projects the input into multiple heads via query (Q), key (K), and value (V) layers (O). Each head independently learns dependencies through scaled dot-product attention. The outputs from all heads are concatenated and merged by a linear layer (2). The MLP then processes the MHA output using position-wise transformations. It first projects the MHA output to a higher dimensional space (3), applies a GeLU (or ReLU) non- -oneeiiear,,y, and then projects back to its original dimension (4). ). As a result, given the block input X₁, the output of the i-th Transformer block can be formulated as:
X, (IN((L)(X)))) + MTP (TN((, + (LN((ILXXX))))
Here, X, is progressively refined by the MHA and MIP, supported by residual connections 12 and layer normalization (LN) [33. These residual connections and layer normalizations are essential

---

Broadcast
All-Reduce
Aggregate
GPUL
GPUI
MHA
MHA
MLP
O MLP
GPU2
GPU2
MLP
MLP
Pre-LN Transformer w/TP
(b) FAL w/TP
Figure 2: Forward pass during Transformer training with tensor parallelism across 2 GPUs. (a)
Standard training involves Broadcast, All-Reduce, and Aggregate communication. (b) FAL training requires only Broadcast and Aggregate, reducing communication overhead.
for stabilizing gradients and preventing rank collapse 14. Pre-LN design forms the basis for many
large-scale models such as GPT (2] and LLaMA [3] In Sec. we build upon this baseline to propose
structural modifications for efficiency and model quality.
Tensor Parallelism for Transformer. Fig. (a) illustrates tensor parallelism (TP) 4.. which is one of the widely used parallel training methods. TP, the parameters the transformer are sliced either along the row (input) or column (output) dimensions and distributed across GPUs.
The input is first broadcasted across GPUs and then fed into the sliced MHA. The outputs of the sliced MHA are then summed up together to be passed to the subsequent MLP. This results in all reluce communication between GPUs to fully assemble the MHA output before passing it the MLP. After that, the assembled MHA output is fed into the sliced MIP. The outputs of the sliced MLP also need to summed up together to passed to the next transformer block - this requires an Prregate-broaccst aggregate- broadcast communication.
Unlike data parallelism [1] and pipeline parallelism [6] which require moving model states or introducing bubble periods, TP can achieve high computational occupancy when the model size is much larger than the activation footprint 7 However, it still suffers from substantial communication overhead as the number of GPUs increases or the interconnect speed slows down [S) [7]. Since TP partitions MHA and MLP parameters across GPUs, each transformer block requires two all-reduces in both forward and backward passes, creating significant overhead. To mitigate this, we focus on reducing the necessity of all-reduce within each transformer block.
Motivation
crucial principle in typical transformer architectures that the MLP must always receive the most recent MHA output within each block. This incurs all-reduce communication in TP affecting the overall training time. Here, given that the residual path in a specifie block already accumulates the attention outputs of all preceding blocks, we question whether the MLP truly needs the most recent MHA output. To answer the question, we conduct several analyses on GPT-2, using four linguistic
datasets WikiText-2 [8] PTB 19 BookCorpus 20 and CC-News 21
3.1 MHA-MLP Connections Can Be Reconfigured
To examine whether the direct MHA-MLP connection is strictly necessary, we first measure how much a single MHA output has impact the MLP input within block. Specifically, quantify the feature representation changes between the MHA output and MLP input by comparing their CKA similarity [22 similar practice used in many prior works quantify the feature representation
changes across activations in the model 23 24
Fig. 3 (a) shows the CKA similarity scores for MHA outputs, MLP inputs (Residual + MHA), and MLP outputs across adjacent blocks. Although the MHA output, which largely varies across the
A comparison of training times among data, pipeline, and tensor parallelism is provided in Apdx. B

---

MHA output
Residual
MHA output BookCorpus
MLP output
Wikitext-2
PTB
CC-News
Original
All MHA
All Connect
3600
Similarity
Perplxxity
0.8
2400
0.4
1200
0
3
10
11
Wikitext-2
BookCorpus CC News
Transformer Block Index
(a) Similarity analysis
(b) Connection ablation
Figure 3: Motivation for reconfiguring MHA-MLP connections. (a) CKA similarity analysis across successive GPT-2 blocks. The x-axis indicates the block index. and the y-axis shows the similarity between consecutive MHA outputs, Residual+MMHA outputs (i.c., MLP inputs), and MLP outputs. (b) Connection ablation results measured by perplexity. Original denotes the unaltered model, All MHA removes all MHA layer, and All Connect removes all direct MHA-MLP connections.
Gradient Magnitude
Wikitext-2
PTB
BookCorpus
CC-News
Normalized
1.2
Normalized
Perplexity
80
0.8
40
0.4
9
10
11
12
8
9
10
11
12
Trunsformer Block Index
Transformer Block Index
(Gridient Magnitude of MHA Output)
(where MHA layer omitted)
(a) Gradient magnitude analysis
(b) Layer ablation
Figure 4: The crucial role of the first attentions. (a) Normalized gradient magnitude of the MHA outputs across Transformer blocks in GPT-2 for different datasets. The x-axis represents the block index. (b) Layer ablation results measured by perplexity. The x-axis indicates the index of the transformer block from which the MHA is omitted.
blocks, is added into the MLP input, the MLP input does not vary much across the blocks. This suggests that single MHA output change has limited impact on the MLP inputs since the residual connection already accumulates attention signals from the previous blocks. Note that the MLP still produces different outputs, meaning that information can be generated even if the inputs are similar. This implies that the MLP may not always require the most recent MHA output (i.e., the output of the MHA within the same block).
Given the observation that MHA output and MIP input within the same block consistently yields similar representations, hypothesize that the MHA-MLP connections be selectively skipped or reconfigured for efficiency gain high representational similarity of intermediate features often leaves potential for lightweight modifications or pruning [23]. To validate this hypothesis, remove either entire MHAs or only their connections to the MLP and quantify the impact of the removals.
Fig. (b) illustrates two scenarios: removing all MHAs (All MIIA) versus removing all MHA- MLP connections (All Connect). As expected, removing All MHA severely degrades model quality, since attention can be utilized all. In contrast, removing All Connect recovers much of the lost performance compared to removing the entire MHAs. Although skipping connections less harmful than removing entire layers, the performance loss remains significant, indicating the need for alternative signals for ML.P
3.2 First Attention is Key
To investigate alternative attention signals for MLP, measure the gradient magnitude of the MHA outputs across all blocks to find a crucial attention output with high impact on final predictions. of the widely used method to measure which input features the
Note gradient magnitude is also model is focusing on as an importance score 255
Fig. shows that the first MHA output consistently exhibits the highest gradient magnitude, indicating that perturbations in the earliest attention result have disproportionately large impact on the final prediction. We further confirm this by measuring the perplexity after omitting the

---

MHA from individual transformer blocks. As shown in Fig. 3 (b), removing the first attention causes a far larger perplexity increase than removing later layers, verifying the crucial role of the first attention in language modeling. These findings align with the well-known psychological phenomenon of the primacy effect [26), commonly summarized as "first impressions matter." The primacy effect of the first attention is not limited to a specific model architecture previous works also identified the prominent impact of the first attention layer across various attention mechanisms and tasks [27 28 29)
Our findings suggest that leveraging the first attention more effectively can compensate for the information loss caused by skipping direct MHA-MLP connections
4 FAL: Harnessing First Attention for Enhanced Efficiency
Our analyses in Sec. indiccaee 1) although MHA outputs can be bypassed, alternative signals for MLP are needed (Sec. (.1). and 2) the first MHA output could be the key to bridging the performance gap (Sec. (.2c.3.2). Building these findings, we propose FAL (First Attentions Last), a novel transformer architecture designed to streamline MHA-MLP connections using the first MHA. We emphasize three key considerations to balance efficiency and model quality:
Motivated by dual process theory 30 which suggests that revisiting the initial judgement can improve accuracy, FAL carefully retains (and dwells on) the crucial first attention signal in the subsequent blocks, by distinctly treating the first block as specialized "preparation" stage.
Retaining Crucial First-Attention Signal:
Maintaining the Transformer Structure: Apart from rerouting the MLP input, FAL retains the conventional ordering of sub-layers, residual paths, and multi head attention components. Our design thus preserves the proven benefits of standard residual connections and MHA while enabling more careful use of the first attention output.
Minimizing Overhead: The reconfiguration must reduce data communication in multi-GPU training while avoiding excessive memory usage computations.
4.1
FAL Architecture
Fig. (b) illustrates the core design of FAL. FAL revisits the first attention output in all subsequent blocks, reducing the need to rely on the most recent attention output at each block. For reuse, we apply layer normalization (LN) to the first attention output only once in the first block, thereby maintaining a stable scale without incurring repeated LN overhead. Because the first attention output passed through every block, FAL. acts akin to a dense skip connection 31 for this highly influential first-stage signal.
The MLP now receives LN(X,) LN (LN(X,))) instead of LN(X; + MHA,
By adding these two LN outputs together, LN(MHA, (MHA, (LN(X1))) and LN(X,). our design makes the LN affine parameters learn the relative weight of each component. F Meanwhile, each block still computes its own MHA (i.e., MHA,), which remains in the residual connection at the output of the block incorporate newly captured information. As a result, the ith FAL block output formulated as:
X; + (LA((L)(X))) + MLP, (LN(()) + (LNM,,,,
4.2 Effectiveness of FAL
Reducing Communication in Multi-GPU Parallelism. With reconfigured connections, FAL lowers communication overhead during multi-GPU training. Fig. (b) illustrates FAL training using TP, which fuses the all-reduce operation for MHA into that for ML.. By connecting the first MHA, instead of the latest one, to the MLP, the MLP no longer requires the all-reduce output. As result, the outputs of the latest MHA and the MLP can be added locally on each GPU. This halves the communication overhead, significantly improving the training time performance.
"Additional analyses (Apdx. C consistently support (1) the feasibility of bypassing blo-k MHA outputs, and (2) the pivotal role of the first MHA output across transformer variants and domains.
'To avoid storing extra activations, we reposition the first layer's LN from the MLP input to the MHA output. This allows later blocks to reuse the normalized first- attention output without recomputation or extra memory, as LN result is already cached during the forward pass for backpropagation.
5

---

Single-GPU Acceleration. FAL also speeds up single-GPU training. When executing standard transformer block on a single GPU, the MLP must wait until the MHA finishes its operations. This se- quential process can lead to sub-optimal utilization of computational resources during training (32).
Kernel Sequence
Original
Muttelape Peceptne
FAL
Mull-HHad Altanion
Improved Paralletion
Mali- syge Pervesbood
Figure 5: Parallel kernel launch of FAL.
As depicted in Figure 3 because FAL modules (i.e., MHA and MLP) have no dependencies, they can be executed in parallel accelerating single-GPU training For example, while memory- intensive operations (e.g., element-wist operations) of a module are executed, compute-intensic operations (e.g., matrix multiplications) of the other module can be simultaneously executed. This concurrency increases both computational and memory throughput, reducing overall training time.
Note, although we focus the training efficiency in this paper, the benefits of FAL. also be applied to multi-GPU inference scenarios - TPP is widely used even for the multi-GPU execution of inference due to its decent computation throughput and memory efficiency (7]. The detailed results of FAL's inference acceleration are presented in Apdx. D.3
FAL+: Harnessing First Attention for Enhanced Quality
Our observations in Sec. 3.2 highlight the crucial role of the first MHA in language modeling. In FAL, exploit the first MHA output reduce overhead by replacing subsequent direct MHA-MLP connections. In this section, we introduce FAL+, which instead augments the original MHA-MLP connections with the first MHA output further enhance model quality.
Fig. D (c) illustrates the core design of FAL+ Rather than removing any MHA-MLP connections, each transformer block integrates the first attention signal MHA (LN(X1)) alongside its original connection. FAL+ appends an additional LN on each block for the first attention signal, allowing the LN affine parameters to control how much of the first attention is utilized.
As we show Sec. 6.4
FAL+ consistently achieves lower perplexity than the baseline, demonstrating that the first MHA output can be leveraged effectively even when the primary goal is to improve model quality rather than training speed.
6 Evaluation Results and Analysis
6.1 Experimental Setup
We conduct our experiments on variety of hardware, datasets, and models with various scales, as briefly summarized below (further details in Apdx. AA).
Hardware: In order to comprchensively evaluate our approach across diverse GPU architectures
and scales, we conduct experiments on multi-GPU configurations (2-8 GPUs) with RTX 3090 and H200 devices connected via PCIe or NVLink, and on single-GPU setups with RTX 3090,
RTX 4090, and RTX A6000.
Baselines: We compare FAL and FAL+ with a standard transformer- language model GPT-2 and larger GPT variants FAL and FAL+ are implemented atop the baselines. We also compare FAL and FAL+ with parallel configuration considered in 33
34 35
where MHA and ML.P
modules are executed simultaneously using the same input, order to validate that reusing the first MHA output shows similar parallelism improvements with the parallel configuration while even enhancing the model quality.
In multi-GPU scenarios, we further compare our proposed architectures with lossy communication time reduction methods, such as gradient quantization (36] and low rank approximation (37.
Datasets: We pre-train the models on Open Web'Tet corpus (38), publicly available counterpart
to GPT-2's WebText. For scalability analysis, we use the Pile dataset [39]. Zero-shot performance is evaluated on language understanding tasks using the SuperGLUE benchmark suite [8]
'In practice, FAL enables overlapped execution of MHA and MLP on separate CUDA streams. Such overlap allows the warp scheduler to better hide latency when one warp stalls on a memory operation, another ready warp in different stream can be scheduled without delay.

---

Normalized Time ferr
GPT2
FAL
Computation
Communication
0.5
774M
1.58
2.58
8.38
774M
1.58
2,58
8.38
774M
1.38
2.58
6.3
774M
.B8
774M
1.58
2GPU
4GPU
8GPU
2GPU
4GPU
11200
w/NVLink
RTX3090
w/PCle
Figure 6: Normalized Multi-GPU Training Time of GPT2 and FAL.
6.2 Multi-GPU Performance
Fig. shows the normalized training time of GPT-2 and FAL on multi-GPUs (i.e., H200 with NVLink and RTX 3090 with PCIe) with model sizes from 774M 8.3B. For high-end GPUs with NVLink, FAL improves the training time performance by 13.2% on average (up to 20.1% 1) compared to GPT-2. In the case of GPUs, communication overhead relatively low due the high bandwidth of NVLink. In this case, FAL's performance improvement mainly comes from the single-GPU acceleration. As the number of GPUs increases for training larger models - the communication overhead substantially increases. In such cases, FAL further improves performance (by 18.7% on average) by reducing all-reduce communication within each transformer block.
In typical multi-GPU servers with PCIe (instead of NVLink) 40) the communication overhead becomes more pronounced accounting for up 80.6% of training time on 4 GPUs. In such cases, FAL provides even greater benefit it improves training time performance by 36.6% on average (up to 43.1% compared GPT-2. These results demonstrate that FAL is not only effective in high-performance setting, but also highly beneficial in typical PCIe-based settings, reinforcing its scalability and practicality across a variety of deployment settings.
Comparison with Lossy Communication Reduction Meth- ods. 7 shows the training time breakdown and perplex- ity of GPT-2, FAL. and two gradient compression methods (quantization (36] and low-rank approximation (37) denoted as Grad-Q and Grad-LR, respectively) on a 2-GPU PCIe setup with WebTee. While the compression techniques sub- stantially reduce communication time (by 37.8% on average), they significantly degrade model quality. On the other hand, FAL reduces much more communication time overhead (by 49.4%) compared to the compression techniques without com- promising the model quality it even reduces the perplexity compared GPT-2. This result demonstrates that FAL can strike much better performance-accural trade-off point com- pared the prior communication reduction techniques the connection reconfiguration.
BWD (DeXComp
PPL
Comm
1.6
Normalized Timelter
1.2
Candaaiion PPL
0.4
GPT-2
GralQQ God-LR
FAL
Figure 7: Validation perplexity and training time breakdown of GPT- 2, gradient compression meth- ods, and FAL. (FWD: forward, BWD: backward, Comm: com- munication, (De)Comp: compres- sion/decompression time)
6.3 Single-GPU Performance
Fig. 8 (a) shows the normalized throughput (tokens per second) GPT-2 and FAL on single GPUs
(RTX 3090, RTX4090, and RTX A6000); the throughput is normalized to that of GPT-2. As shown in Fig.1 (a), FAL improves single-GPU throughput by 1.08x on average (up to 1.18x) compared to GPT-2. This is because FAL enables overlapped execution of MHA and MLP in each block better utilizing resources - Fig. S (b) shows that FAL improves SM utilization, warp occupancy, tensor core usage, and memory bandwidth by up to 8.2%, 45.9%. 13.9%, and 18.4%. respectively, on RTX-3090.
Note FAL typically shows better single GPU throughput when FlashAttention is adopted. This is because FlashAttention increases the computational intensity of attentions with kernel fusion, giving more opportunity for FAL to overlap the computation-intens. intensive operations in MHA and MLP leading to better resource utilization. Although MLP's large matrix multiplies are compute-heavy. each GEMM begins and ends with global- loads and stores. These boundary memory transactions introduce unavoidable stalls, even if the core GEMM
operations and memory-

---

GPT2
FAL.
Narmalized Throughpuu
1.2
Unlizztion
Warp Occupmncy
100%
60%
1.1
Percentage (%)
90%
Percertang (%)
40%
Onggiall Flash Atta
Original
Flash Atmn
Tensor Core Usage
Momory Bandwiddh
0.9
30%
60%
b-1 b--
b-1
b-2
b-1
b-2
b-]
b-
b-1
b-1
b-1 b-1
Original
Original
Original
70%
40%
RTX 3090
RTX 4090
RTX A6000
10%
205
Original
Flush Altn
Original Flash Alin
(a) Single-GPU throughput
(b) Single-GPU Profiling
Figure 8: Single-GPU throughput comparison between GPT2 and FAL. (a) Normalized through- put (tokens/sec). (b) Analysis of throughput gains on RTX 3090.
Table 1: Openwehtext validation perplexity, training time, and SuperGI.UE. zero-shot results.
Openwebtext (
SuperGLUE (CB, Record: Fl
Others: Accuracy)
PPI. Timee
BoolQ
CR
MultiCC 45
ReCoRDD
WIC
Avg.
GPT-2 774M
17.75/ 13.2d
55.7
19.4
$4.0)
52.3
57.4
54.2
49.8
45.2
48.5
Paralle
17.80/8.6d
50.0
19.4
58.0
53.8
48.6
51.6
49.1
36.5
15.9
17.55/8.6d
54.5
48.5
FAL+
17.24/ 13.24
51.8
21.1
58.0
55.7
56.2
51.3
51.3
48.1
49.2
GPT-2 1.5
14.72/24.16
$8.0
24.1
65.0
57.2
78.4
53.1
50.0
40.4
53.3
FAL
14.23 16.1/
58.1
21.6
72.0
57.2
78.2
$4.2
49.2
64.4
56.9
14.12/24.20
79.0
51.0
is compute- bound. FlashAttention's higher arithmetic intensity in the attention block lengthens its compute phase relative its memory phase, creating more opportunities to hide MLP's boundary stalls behind attention computation.
6.4 Model Quality and Training Efficiency
Training from Scratch. To validate the model quality of FAL and FAL+, measure the perplexity and end-to-end training time, respectively, while training GPT-2, parallel configuration (denoted as Parallel), FAL, and FAL+ with Open WebText dataset 4-GPU PCIe setup.
Table (left) reports the perplexity and total training time for the 774M (36 layer) and 1.5B (48 layer) scales. Both FAL and Parallel improve the training time by 34%, average. compared to GPT-2. However, Parallel degrades model quality compared to GPT-2, since it discards MHA-MLP connections without providing any alternative features. On the other hand, FAL even improves model quality - lowers perplexity by 0.2 and 0.49, compared to GPT-2 774M and 1.5B, respectively. This result demonstrates that reconnecting the first MHA output with LN is not just alternative signal the removed connections, but a deliberate reuse of the crucial early representation which leads to better understanding of the input.
FAL+ achieves even lower perplexity compared to GPT-2, by augmenting the MHA-MI.P connections with the first MHA output its training time is thus almost same with the baseline. Here. the perplexity improvements of FAL and FAL+ get larger as model scale increases this is mainly because the advantage of revisiting the first attention becomes increasingly effective in deeper models (see Scalability Analysis for more details).
Scalability, To validate the scalability of FAL. and FAL+, we compare token ingestion efficiency of the Pre-LN Transformer, FAL, and FAL+ as their depth increases from to 60 layers. Fig. p shows the loss curves under fast training conditions inspired by Cramming [50]. In the carly stages of training, both FAL and FAL+ reduce MLM loss more rapidly than the Pre-LN baseline. In case of 36-layer scale, all models end up show similar loss values.
As the model depth increases, both FAL and FAL+ converge lower MLM loss values than the Pre- LN baseline - FAL+ usually exhibits lower MLM loss values compared to FAL. In case of standard Pre-LN transformers, residual connections cannot fully preserve early-layer information, which gradually dilured activations accumulate through depth (9. This limits deeper architectures from revisiting the high-impact first attention signal that strongly influences the model's final predictions. In contrast, FAL and FAL+ repeatedly reintroduce the high-impact first attention output across depth,

---

Pre-LN
FAL
FAL
Blocks
of Blocks 48
Blocks
4.5
4.5
MLM Loss
3.5
200M
400M
600M
S00M
18
200M
400M
C00M
800M
1B
200M
400M
600M
H0OM
# of Ingested Tokens
Figure 9: Loss comparison with increasing number of blocks across Pre-LN architecture, FAL, and FAL+.
allowing later blocks adaptively re-weight this signal. Moreover, the degradation from removing direct MHA-MLP connections diminishes in deeper models, while the primacy of the first attention persists (as we demonstrate in Apdx. C). This validates the benefit of the replacement strategy of FAL in deeper architectures.
Generalizability. To further validate the generalizability of FAL and FAL+. we evaluate their zero-shot performance on language understanding tasks using SuperGLUE benchmark. Table [ (right) shows the zero-shot results, where CB and ReCoRD are evaluated using Fl score, while the remaining tasks use accuracy. FAL largely preserves the language modeling capabilities of the standard Pre-LN architecture. In the case of 774M scale, FAL achieves the same average score as GPT-2 across SuperGLUE tasks. On the other hand, in the case of 1.5B scale, FAL surpasses GPT-2 due to the increased depth which gives more opportunity to revisit the first attention signal. FAL+ achieves higher average scores than GPT-2 in both scales, by augmenting the MHA-MLP connections with the first MHA output. Our motivational analyses vision and multimodal tasks in Apdx. C also demonstrate the feasibility of bypassing per-block MHA outputs using the pivotal first MHA output. Hence, we believe FAL and FAL+ can further be generalized to other tasks.
We also validate FAL and FAL+ on various transformer variants, such as LLaMA (with Grouped Query Attention, GQA) and Switch Transformer (with Mixture of Experts, MoE). Both FAL and FAL+ consistently yield lower loss curves than the baseline GQA and MoE architectures, confirming that our connection reconfiguration is broadly applicable the detailed results and training configurations are provided in Apdx.
Downstream Robustness. We empirically amine how injecting the first attention signal during instruction tuning affects the trade-off between stability and adaptation. We fine-tune both GPT-2 1.5B and FAL+ 1.5B on the Alpaca instruction tuning dataset (53 using four learn- ing rates (le-5 to le-2).
Table 2: Instruction tuning robustness Al- paca. (stability vs. adaptation)
Model
LR
Val PPL
Trained PPL
GPT2 1.5B
lc-5
0.01
30.93
1c-4
0.00
30.92
le-3
0.01
28.10
Table Table[rreport trained perplexity (Trained PPL) on Alpaca, reflecting adaptation, and validation perplexity degradation (A Val PPL) on Open- WebText, reflecting forgetting. FAL+ consis- tently preserves pretraining knowledge better than GPT-2. Across all learning rates (LR), FAL+ shows lower or equal A Val PPL on Open-
le-2
1.39
4.83
FAL+ 1,5B
le-5
0.00
27.97
le-4
0.00
27.98
le-3
-0.01
27.22
le-2
0.61
5.76
'Although FAL+ exhibits lower perplexity compared FAL during pretraining. its SuperGLUE lower than FAL. This discrepancy suggests that domain mismatch between the Open WebText corpus and SuperGLUE tasks may limit the direct transferability of perplexity improvements on the pretraining corpus to downstream performance 52 However, we expect this gap can be mitigated through continual pretraining more diverse domain sources [11].

---

WebText compared to GPT-2, indicating stronger stability. At Lt= FAL+ even shows slight improvement validation PPL (-0.01), suggesting positive regularization rather than forgetting.
F also achieves better adaptation when forgetting is minimal. In Table while maintaining zero negative APPI., FAL+ attains the lowest trained PPI. (27.22) at LR=1e-3,, demonstrating high adaptability with minimal forgetting. In contrast, GPT-2 reaches a lower trained PPL (4.83) at LR=le- 2 only by severely compromising original domain knowledge (A Val PPL 1.39). These results show that reusing the impactful first attention in FAL+ enables robust adaptation while mitigating catastrophic forgetting [54), ensuring that the improved model quality (Table [ is not achieved the cost of stability or downstream adaptability.
7 Related Work
Training large-scale neural networks requires parallel methods to handle high computation and memory demands. Data Parallelism (DP), Pipeline Parallelism (PP), and Tensor Parallelism (TP) are the primary approaches. In DP, each GPU processes different data batches with identical model replicas and synchronizes gradients, but scaling to large models incurs heavy memory and communication costs. PP [16] partitions (sub layers into sequential stages across GPUs, introducing pipeline bubbles that slow training. TP [4] distributes model (sub)layers across GPUs to avoid bubbles, but communication still remains bottleneck. FAL tackles this by exploiting the first attention to reduce communication while preserving information. Some works 5 [55] combine TP and PP with DP for efficient GPU training. FAL can further boost these combinations by lowering the communication cost introduced by TP.
Parallel Training Methods:
Communication Reduction in Parallel Training: Overlapping communication with computation can hide communication delays [6], but its effectiveness depends on how much communication can be
overlapped with the computation [56]. FAL offers more opportunity overlap communication with computation by reducing the communication frequency. Gradient compression using quantization [36] or low-rank approximation 37 reduces data exchange, but the (de)commression overhead can negate time savings given TP's short communication intervals.
lower TP communication 7 by using the same input to the MHA and MIP, yet degrade quality on linguistic tasks. In contrast, FAL. reconfigures sequential connections using the first MHA output, thereby cutting communication overhead while improving model quality.
Parallel configurations within transformer blocks 33 34 35
8 Conclusion
Training large transformer models over multiple GPUs often incurs substantial communication overhead. To alleviate the communication overhead, we propose FAL (First Attentions Last), novel architecture that leverages the output of the high-impact first MHA to streamline MHA-MLP connections. Leveraging the first attention more effectively, FAL removes expensive all-reduce communication within each block and enables parallel execution of MHA and MLP on a single GPU without compromising the model quality. We also introduce FAL+, a variant that leverages the first attention output to pursue additional quality gains, underscoring the flexibility of our connection-level approach. In our evaluation across various linguistic tasks and hardware configurations, FAL shows up to 44% of training time reduction, compared to the baselines of GPT architecture, while improving model quality as the depth increases. Furthermore, FAL+ achieves even better model quality with connection augmentation. We believe this connection-level perspective opens new avenues for refining transformer architectures, both in terms of communication efficiency and model quality.
Limitation and Future Work. FAL slightly degrades accuracy for small models (0.3%, Apdx. E.2 because shallow networks provide fewer layers to accumulate attention signals, making the replace- ment of block-level connections less stable. FAL+, however, augments rather than replaces these connections, allowing later blocks to leverage both the current attention and the first attention, which avoids information loss and even yields slight accuracy gain (0. 14%). Building upon this, in- corporating a commeciion-reconiguraa Architecture Search [57] or dynamically injecting the first attention through gating (58 presents an interesting direction for future research.
strategy centered on the first attention output into Neural
10

---

Acknowledgements
This work was supported in part by National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (RS-2023-00212711 I, RS-2025-24534857, and RS-2025-25434742 Institute of Information & Communications Technology Planning & Evaluation (LITP)- - ITRC (Infor- mation Technology Research Center) grant funded by the MSIT and ICT Creative Consilience Program through IITP grant funded by the MSIT (ITTP-2025-4) 1-25-RS-2000 2020- I1201819). We also thank the members of Computer Architecture & Systems research lab in Korea University for their useful comments and discussions, as well as the anonymous reviewers for their helpful feedback.
-SS-RS--222-000260991
References
[1] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint ar arXiv:2001.08361, 2020.
[2] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidr, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint (rXXi:2303.08774 2023.
[3] Hugo Touvron, Thibaut Lavril, Gautier Izacard. Xavier Martinet, Marie-Anm Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 2023.
[4] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models using model parallelism. arXiv preprint 1rXvv:190.00803,, 2019.
[5] Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, al. Efficient large-scale language model training on gpu clusters using megatron-Im. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1-15, 2021.
[6] Shibo Wang, Jinliang Wei, Amit Sabne, Andy Davis, Berkin Ilbeyi, Blake Hechtman, Dehao Chen, Karthik Srinivasa Murthy, Marcello Maggioni, Qiao Zhang, et al. Overlap communication with dependent computation via decomposition in large deep learning models. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1, pages 93-106, 2022.
[7] Reiner Pope. Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently scaling transformer inference. Proceedings of Machine Learning and Systems, 5:606-624, 2023.
[8] Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Superglue: A stickier benchmark for general-purpose language understanding systems. Advances in neural information processing systems, 32, 2019.
[9] Ruibin Xiong, Yunchang Yang. Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization the transformer architecture In International conference on machine learning, pages 10524-10533, PMLR, 2020.
[10] Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. arXiv preprint arXiv:2004.08244 2020.
[11] Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse. Advances in Neural Information Processing Systems, 35:27198 15:27198-27211, 2022.
[12] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
11

---

[13] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450, 2016.
[14] Yihe Dong, Jean -aptiste Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. In International Conference on Machine Learning, pages 2793-20 7793-2803. PMLR, 2021.
[15] Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthicu Devin, Mark Mao, Marc' aurelio Ranzato, Andrew Senior, Paul Tucker, Ke Yang. al. Large scale distributed deep networks. Advances in neural information processing systems, 25, 2012.
[16] Yanping Huang. Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, Hy- oukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems, 32, 2019.
[17] Shenggui Li, Hongxin Liu, Zhengda Bian, Jiarui Fang, Haichen Huang, Yuliang Liu, Boxiang Wang, and Yang You. Colossal-ai: A unified deep learning system for large-scale parallel training. In Proceedings of the 52nd International Conference on Parallel Processing, pages 766-775, 2023.
[18] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.
[19] Mitchell P. Marcus, Beatrice Santorini, and Mary Ann Marcinkiewiez. Building large annotated corpus of English: The Penn Treebank. Computational Linguistics, 19(2)313-3330, 1993.
[20] Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In The IEEE International Conference on Computer Vision (ICCV), December 2015.
[21] Felix Hamborg, Norman Meuschke, Corinna Breitinger, and Bela Gipp. news-please: A generic news crawler and extractor. Proceedings of the 15th International Symposium of Information Science, pages 218-223, March 2017.
[22] Simon Kormblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited. In International conference on machine learning, pages 1519.3529. PMLR, 2019.
[23] Thao Nguyen, Maithra Raghu, and Simon Kornblith. Do wide and deep networks learn the same things? uncovering how neural network representations vary with width and depth. arXiv preprint arXiv:2010.15322 2020.
[24] Hyegang Son, Yonglak Son, Changhoon Kim, and Young Geun Kim. Not all adapters matter: Selective adapter freezing for memory-efficiee fine-tuning of language models. In Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computa- tional Linguistics: Human Language Technologies (Volume 1: Long Papers). Association for Computational Linguistics, April 2025.
[25] Karen Simonyan, Andrea Vedaldi, and Andrew Zisserman. Deep inside convolutional networks: Visualising image classification models and saliency maps. arXiv preprint arXiv:1312.6034 2013.
[26] Solomon E Asch. Forming impressions of personality. The journal of abnormal and social psychology, 41(3):258, 1946.
[27] Maximiliana Behnke and Kenneth Heafield. Losing heads in the lottery: Pruning transformer. In The 2020 Conference on Empirical Methods in Natural Language Processing, pages 2664-2674. Association for Computational Linguistics (ACL), 2020.
12

---

[28] Xingwu Chen, Lei Zhao, and Difan Zou. How transformers utilize multi-head attention in in-context learning? case study on sparse linear regression. arXiv preprint arXiv:2408.045:2 04533.... 2024.
[29] Yang Zhang, Yanfei Dong, and Kenji Kawaguchi. Investigating layer importance large language models. arXiv preprint arXiv:2409. 14381, 2024.
[30] J St BT Evans. Dual processes in reasoning? Cogmition, 3(2):141-154, 1974.
[31] Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger. Densely connected convolutional networks. In Proceedings the IEEE conference on computer vision and pattern recognition, pages 4700-4708, 2017.
[32] Andrei Ivanov, Nikoli Dryden, Tal Ben-Nun, Shigang Li, and Torsten Hoefler. Data movement is all you need: A case study on optimizing transformers. Proceedings of Machine Learning and Systems, 3:711-732, 2021.
[33] Ben Wang and Aran Komatsuzaki. GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model. https: //github// com www./iiigooollo//mmeeeioioooo transf -jax May 2021.
[34] Bobby He and Thomas Hofmann. Simplifying transformer blocks. arXiv preprint
arXiv:2311.01906. 2023.
[35] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1- 113, 2023.
[36] Dan Alistarh, Demjan Grubic, Jerry Li, Ryota Tomioka, and Milan Vojnovic. Qsgd: Communication-efficiien sgd via gradient quantization and encoding. Advances in neural information processing systems, 30, 2017.
[37] Thijs Vogels, Sai Praneeth Karimireddy, and Martin Jaggi. Powersgd: Practical low-rank gradient compression for distributed optimization. Advances in Neural Information Processing Systems, 32, 2019.
[38] Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex. Openwebtext corpus. http://Skyliono077 Satyyliioo00-githubb. 1o//penneeTTextCorput 2019.
[39] Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.0000. 2020.
[40] Heehoon Kim, Junyeol Ryu, and Jaejin Lee. Tecl: Discovering better communication paths for pcie gpu clusters. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3, pages 999-1015, 2024.
[41] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with fo-awareness. Advances in Neural Information Processing Systems, 35:16344-16333 2022.
[42] Christopher Clark, Kenton Lee, Ming- Ming-Wei Chang, Tom Kwiatkowski. Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. arXiv preprini arXiv: /905. 10044, 2019.
[43] Marie-Catherine De Marneffe, Mandy Simons, and Judith Tonhauser. The commitmentbank: Investigating projection in naturally occurring discourse. In proceedings of Sinn und Bedeutung, volume 23, pages 107-124, 2019.
[44] Melissa Roemmele, Cosmin Adrian Bejan, and Andrew S Gordon. Choice of plausible alterna- tives: An evaluation of commonsense causal reasoning. In AAAI spring symposium: logical formalizations of commonsense reasoning, pages 90-95, 2011.
13

---

[45] Daniel Khashabi, Snigdha Chaturvedi, Michael Roth, Shyam Upadhyay, and Dan Roth. Looking beyond the surface: A challenge set for reading comprehension over multiple sentences. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volume I (Long Papers), pages 252-262, 2018.
[46] Sheng Zhang. Xiaodong Liu, Jingjing Liu, Jianfeng Gao, Kevin Duh, and Benjamin Van Durme. Record: Bridging the gap between human and machine commonsense reading comprehension. arXiv preprint arXiv: 12885, 2018.
[47] Ido Dagan, Oren Glickman, and Bernardo Magnini. The pascal recognising textual entailment challenge. In Machine learning challenges workshop, pages 177-190. Springer, 2005.
[48] Mohammad Taher Pilehvar and Jose Camacho-Collados. Wie: the word-in-context dataset for evaluating context-sensitive meaning representations. arXiv preprint arXiv: arXiv: 2018.
[49] Hector J Levesque, Ernest Davis, and Leora Morgenstem. The winograd schema challenge. KR, 2012:13th, 2012.
[50] Jonas Geiping and Tom Goldstein. Cramming: Training a language model on a single gpu in one day. International Conference on Machine Learning, pages 11117-11143, PMLR, 2023.
[51] Suchin Gururangan, Ana Marasović, Swabha Swayamdipta, Kyle Lo. Iz Beltagy, Doug Downey, and Noah Smith. Don't stop pretraining: Adapt language models to domains and tasks. arXiv preprint arXiv:2004. 10964, 2020.
[52] Yi Tay, Mostafa Dehghani, Jinfeng Rao, William Fedus, Samira Abnar, Hyung Won Chung, Sharan Narang, Dani Yogatama, Ashish Vaswani, and Donald Metzler. Scale efficiently: Insights from pre-training and fine fine-tuning transformers. arXiv preprint arX.v:2109. 10686, 2021.
[53] Rohan Taori, Ishaan Gulrajani, Tianyi Zhang. Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford alpaca: An instruction-followinn Ilama model. https://g/ thub. com
www/tatsu- tatsttttulaaa latssuu-lb/stannf ord alpaca 2023.
[54] German I Parisi, Ronald Kemker, Jose Part, Christopher Kanan, and Stefan Wermter. Continual lifelong learning with neural networks: A review. Neural networks, 113:54-71, 2019.
[55] Jaeyong Song, Jinkyu Yim, Jaewon Jung, Hongsun Jang, Hyung-Jin Kim, Youngsok Kim, and Jinho Lee. Optimus-ce: Efficient large nlp model training with 3d parallelism aware communication compression. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2, pages 560-573, 2023.
[56] Chang Chen, Xiuhong Li, Qianchao Zhu, Jiangfei Duan, Peng Sun, Xingcheng Zhang, and Chao Yang. Centauri: Enabling efficient scheduling for communication-compuatiii model training via communication partitioning. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3, pages 178-191,2022
overlap in large
[57] Krishna Teja Venkat, Murali Emani, Venkatram Vishwanath, and Arun K Somani. Neural architecture search for transformers: A survey. IEEE Access, 101108344-00412, 2022.
[58] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-3,,20 2022.
[59] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 18):9, 2019.
[60] Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, al. Mixed precision training. arXiv preprint 17X0v:1700 2017.
14

---

[61] NVIDIA Corporation. NVIDIA Nsight Systems. nsight- -yystess 2025. Accessed: 2025-05-02.
https://developor www.:/opprriiiiii.....a.....aa.. nvidia. com/
[62] Ge Yang, Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tuning large neural networks via zero-shot hyperparameter transfer. Advances in Neural Information Processing Systems, 34: 17084- 2021.
[63] Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team, An empirical model of large-batch training. arXiv preprint arXiv:1812.008 2018.
[64] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 2018.
[65] Leslie N Smith and Nicholay Topin. Super-convergence: Very fast training of neural networks using large learning rates. In Artificial intelligence and machine learning for mulli-domain operations applications, volume 11006, pages 369-386. SPIE, 2019.
[66] Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, et al. Pytorch distributed: Experiences on accelerating data parallel training. arXiv preprini arXiv:2006. 15704, 2020.
[67] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimiza- tions toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1-16. IEEE, 2020.
[68] Quan Zhou, Haiquan Wang. Xiaoyan Yu, Cheng Li, Youhui Bai, Feng Yan, and Yinlong Xu. Mpress: Democratizing billion-scale model training on multi-gpu servers via memory-saving inter-operator parallelism. In 2023 IEEE International Symposium on Perforrerforance Computer Architecture (HPCA), pages 556-569, IEEE, 2023.
[69] Soojeong Kim, Gyeong-In Yu, Hojin Park, Sungwoo Cho, Eunji Jeong, Hyeonmin Ha, Sanha Lee, Joo Seong Jeong, and Byung- Syung-Gon Gon Chun. Parallax: Sparsity-awe -ware data parallel training of deep neural networks. In Proceedings of the Fourteenth EuroSys Conference 2019, pages 1-15, 2019.
[70] Yosuke Oyama, Naoya Maruyama, Nikoli Dryden, Erin McCarthy, Peter Harrington, Jan Balewski, Satoshi Matsuoka, Peter Nugent, and Brian Van Essen. The case for strong scaling in deep learning: Training large 3d cnns with hybrid parallelism. IEEE Transactions on Parallel and Distributed Systems, 32(7): 64552 2020.
[71] Alham Fikri Aji and Kenneth Heafield. Sparse communication for distributed gradient descent. 2017.
arXiv preprint
[72] Aaron Harlap, Deepak Narayanan, Amar Phanishayee, Vivek Seshadri, Nikhil Devanur, Greg Ganger, and Phil Gibbons. Pipedream: Fast and efficient pipeline parallel dnn training, arXiv preprint arXiv: I806v 03377, 2018.
[73] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x 16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 2020.
[74] Denis Kocetkov, Raymond Li, Loubna Ben Allal, Jia Li, Chenghao Mou, Carlos Muñoz Fer- randis, Yacine Jernite, Margaret Mitchell, Sean Hughes, Thomas Wolf, Dzmitry Bahdanau, Leandro von Werra, and Harm de Vries. The stack: tb of permissively licensed source code. Preprint, 2022.
[75] Shunyu Yao, Binghui Peng, Christos Papadimitriou, and Karthik Narasimhan. Self-attention networks can process bounded hierarchical languages. arXiv preprint 1111: 2021.
15

---

[76] Rohan Baskar Prabhakar, Hengrui Zhang, and David Wentzlaff. Kraken: Inherently parallel transformers for efficient multi-device inference. Advances in Neural Information Processing Systems, 37:7957-7980, 2024.
[77] Joshua Ainslie, James TeerThorp. Michiel De Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint arXiv:2305.13245 2023.
16

---

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 claims made in the abstract and introduction are clearly reflecting the contribu- tions and scope of this paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contribu- tions made in the paper and important assumptions and limitations. A No or 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 the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We have discussed the limitation of this work in Sec.
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 are encouraged to create a separate "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results to vi- olations of these assumptions (e.g.. independence assumptions, noiseless settings, model sell-specification asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated 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 on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by review- ers as grounds for rejection, a worse outcome might that reviewers discover limitations that aren't it acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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 a complete (and correct) proof?
Answer: [NA]
17

---

Justification:
There is no theoretical result in this paper.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proof' in the paper should be numbered and cross-referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by
formal proof's provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main experi- mental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data provided not)?
Question:
Answer: [Yes]
Justification: The code is included in the supplementary materials. In the paper, the main experi- ment section (Sec. 0 and the section about technical details (Apdx. A provide all information needed to reproduce the results.
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 to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make 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 to accomplish this, but reproducibility can also provided via detailed instructions for how to replicate the results, access hosted model (e.g., in the case of a large language model), releasing of model checkpoint, or other means that are appropriate to the research performed.
example, if the contribution
While NeurIPS does not require releasing code, the conference does require all submissions provide some reasonable avenue for reproducibility, which may depend the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way reproduce the model (e.g., with an open-soure dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe 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.. registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
18

---

Justification:
We included all the important hyperparameters and datasets we used in our experi- We will also open-source our code after the review process.
ments in Apdx.
Guidelines:
The answer NA means that puper does include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (tttps::/nips.cc/pubbll
for more details.
While we encourage the release of code and data, 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 (c.g.. for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/
osoiiiiiiiicciiooooo 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.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments 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) is recommended, but including URLs data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All the information necessary to train the models on Open WebText, perform the comparison using SuperGLUE, replicate results can be found in the paper (Sec. 6 and Apdx.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of 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, or 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: [No]
Justification: We did have the compute resources train several versions of the same model configuration and report perplexity and SuperGLUE results with error bars.
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, confidence intervals, or statistical significance tests, at least for the experiments that support the main
claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar the standard deviation or the standard error of the mean.
19

---

It is OK to report I-sigma error bars, but one should state it. The authors should preferably report 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 symmetric error bars that would yield results that are out of range (e.g. negative error rates).
show in tables or figures
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information the computer resources (type compute workers, memory, time of execution) needed to reproduce the experi- ments?
Answer: [Yes]
Justification: The number of GPU hours necessary to train the models discussed in the paper (Sec. other details of the used computed resources are summarized in Apdx.
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 of compute required for each of the individual experi- mental runs well as 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 it into the paper).
Code of ethies
Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http:///wuuurp....cpppppacccccca
Question:
Answer: [Yes]
Justification:
The Code of Ethics was reviewed carefully for full compliance.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
a The authors should make sure preserve anonymity (e.g., if there is special consideration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
While the paper does present pretrained language models, the accompanying societal impacts and necessary safeguards are different from that of existing language models. As such, we expect that existing mitigation strategies will be applicable to the ideas presented in this paper.
Justification:
Guidelines:
The answer NA means that there is societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no socictal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deploy- ment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
20

---

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 any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement 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 generic algorithm for optimizing neural networks could enable people to 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 to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses no such risks.
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 to usage guidelines or restrictions access the model 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 not 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:
The creators and original owners of assets used in this paper are properly credited, and the license and terms of use are explicitly mentioned and respected, as evidenced by thorough citations.
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, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from 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 use 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 dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's ss creators.
13. New assets
21

---

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [No]
Justification:
This paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their sub- missions 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 is 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
Question: For crowdsoureing experiments and research with human subjects, does the paper include the full of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The experiments in this work did not directly include data from crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution 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 Ethics, workers involved 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 (TRB) approvals or equivalent for research with human subjects
Question: 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 an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: None of the experiments in this work involved human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor 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 the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere the NeurlPS 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.
22

---

Appendix
In this appendix, each section provides the following:
Section A (Technical Appendix): detailed hardware and software configurations used in
experiments, including system specifications and common settings.
Section (Parallel Training Methods); comprehensive background on major distributed training paradigms (Data, Pipeline, and Tensor Parallelism).
Section C (Additional Motivation Analyses): extended motivation analyses across different model scales, datasets, and architectures (GPT-2, ViT-B, LLaMA2-7B, CodeLLaMA-334 and LLaMAA3.2-1BB-issio confirming (1) the feasibility of bypassing per-block MHA outputs, and the pivotal role of the first MHA output across transformer variants and domains.
Section D (Additional Evaluation Results and Analyses): further ablation studies, analyses on information-dilutiin mitigation, and inference acceleration results demonstrating how FAL improves both training and inference efficiency while preserving model stability and quality through effective reuse of the first attention output.
Section E (Evaluation of Generalizability to Transformer Variants): evaluation of FAL and FAL+ on diverse transformer variants (e.g., GQA-based, MoEEbased, and VIT architectures) confirm adaptability across attention mechanisms and modalities.
A Technical Appendix
A.1 Hardware Software
We performed the experiments using PyTorch and Colossal-Al server and public cloud service.
Common Settings
Version of Py PoToch: 2.2.2
Version CUDA: 12.3
Version Colossal-Al: 0.4.0
System-1
Operating system: Ubuntu 20.04.6
CPU: AMD EPYC 7542 32-Core
GPU: NVIDIA RTX 3090 24GB X
Interconnect: PCIe Gen4 x16 (64GB/s)
System-2
Operating system: Ubuntu 20.04.6
CPU: AMD Ryzen Threadripper 3970X 32-Core
GPU: RTX 4090 24GB 2
Interconnect: PCIe Gen4 x16 (64GB/s)
System-3
Operating system: Ubuntu 20.04.6
CPU: AMD Ryzen Threadripper 3970X 32-Core
GPU: NVIDIA RTX A6000 48GB X 2
Interconnect: PCIe Gen4 x16 (64GB/s)
23

---

System-4 (Public Cloud)
Operating system: CentOS 7.9
CPU: Intel Xeon Emerald Rapids (Platinum 8558) 2.10GHz (48-core) / socket
GPU: NVIDIA H200
Interconnect: NVIDIA NVLink (900GB/s)
Tigure[(aa
System:
GPU#: 8
Model: GPT-2 774M
Sequence Length: 1024
Batchsize: 128
Figure 3 (a), (b), (a), (b) Motivation analyses are done with pretrained GPT-2 model on
four different text datasets WikiTeext22 [8]. PTB 19] BookCorpus [20]. and CC-News 21
System:
Model: GPT-2 117M (pretrained: openai-communityggt2-anggg 59 from Hugging Face)
Max Sequence Length: 1024
Figure
We employ FlashAttention 41 and miixdd-precisiie training 60) in all experiments to
maximize tensor core utilization and overall training efficiency. We benchmark using the largest batch size (in powers of two) supported under each training setting.
System: 4, 1
GPU#: [2, 8]
Model: [GPT-2 774M, 1,5B, 2.5B, 8.3B]
Batchsize: System4: 64 (774M, 2GPU), 16 (1.5B, 2GPU), 16 (2.5B, 2GPU), 8 (8.3B, 2GPU), 64 (774M, 4GPU), 32 (1.5B, 4GPU), 32 (2.5B, 4GPU), 16 (8.3B, 4GPU), 128 (774M, 8GPU), 64 (1.5B, 8GPU), 64 (2.5B, 8GPU), 32 (8.3B, 8GPU) Systeml: (774M, 2GPU), 2 (1.5B, 2GPU), 8 (774M, 4GPU), (1.5B, 4GPU)
Sequence Length: 1024
Figure 7
System: 1
GPU#:
Model: GPT-2 774M
Total Batchsize: 32 (used gradient accumulation)
Sequence Length: 1024
Epochs:
Learning rate: 0.0001
Weight decay: 0.001
Dropout: 0
Figure 8 (a) We compare GPT-2 and FAL under both the minimum (1) and maximum batch sizes for each GPU setting. We also evaluate the speedup with and without acceleration techniques specifically, FlashAttention on each GPU configuration.
System: [1, 3]
GPU#: [1]
Model: GPT-2 774M
Sequence Length: 1024
24

---

Figure (b)
We use NVIDIA Nsight Systems [61] to profile GPU performance, including SM utilization, warp occupancy, tensor core usage, and memory bandwidth.
System:
GPU#:
Model: GPT-2 774M
Sequence Length: 1024
Table We train each architecture on Open WexTet [38. an open-source replication of the WebText dataset originally used train GPT-2. The dataset comprises approximately 41.7 GB of text, corresponding to 4 billion tokens. Given our limited computational resources, we use a compute- efficient batch size of 32, which has been shown be sufficient for stable hyperparameter transfer in yp- ua--based training 62 633 To evaluate language understanding performance, we report zero- shot results on the SuperGLUE benchmark [8], which includes BoolQ 42 CB 43), COPA [44).
MultiRC [5), ReCoRD A6). RTE [47). WiC [8]] and WSC 44). No finetuning or additional training was performed on any task. CB and ReCoRD are evaluated using Fl score, while remaining tasks use accuracy.
System:
Epochs:
GPU#: 4
Model: GPT-2 774M, 1.5B
Parallel setting: 2TP/2DP
Total batchsize: (used gradient accumulation)
Sequence Length: 1024
Learning rate: 0.0001
Weight decay: 0.001
dip-grra-norm:
lmbdd-pdrop: 0.1
which demonstrated that scaling laws [1] can be observed even under small-scale, fast- fest-training settings, we compare FAL and FAL+ to the standard pre-LN architecture by stacking transformer blocks with depths of 36 (equivalent to GPT-2 774M), 48 (GPT-2 1,5B), and 60. To evaluate scalability, we stack the transformer blocks of a pre-LN masked language model architecture based on BERT-Large (64.
Figure 9 Motivated by Cramming [50],
Training settings follow the original Cramming paper, including a budget- one-cycle learning rate scheduler [65] and batch size ramp-up for 24 hours I GPU. Models of the same scale are trained under identical system configurations.
System: 1,2
GPU#:
Hidden size: 1024
Intermed size: 4096
Nonlinear: GELU
Max sequence length: 128
Number transformer block: 36, 48, 60
Final Batchsize: 8192 (used gradient accumulation)
Learning rate: 0.0001
Weight decay: 0.001
-oip-grad-norm:
Hidden dropout probability: 0.1
Attention dropout probability: 0.1
Embedding dropout probability: 0.1
25

---

B Parallel Training Methods
DP
GPLO
Made
FWD
BWD
Comm
Moodi
Bubble
Memory Usage
1.2
24
GPU!
PP
Normalized Timefier
18
Memory Usager CPUU
Nask
0.8
Lest
0.6
TP
0.4
GPUI
PP
FAL
Figure 10: Comparison of Data parallclism(DP), Pipeline Parallelism(PP), Tensor Parallelism(TP) We stack GPT-2 blocks until DP can handle them (number of blocks: 42), using Open WebText (max sequence length: 1024) for comparison on two NVIDIA RTX 3090 GPUs connected via PCIe.
To address challenges of training the scale models, employing multiple distributed GPUs along with parallel training methods [66 15 67 68
69 70 has become common practice. These methods encompass a range of parallelism paradigms, including Data Parallelism (DP), Pipeline Parallelism (PP), and Tensor Parallelism (TP). DP duplicates the entire model across multiple distributed GPUs. Each GPU then trains the duplicated model with different data batches and synchronizes the trained gradients for unified updates [1S] 69 71]. While DP is effective for smaller models, it results in significant memory and communication overhead for larger models as each GPU needs retain the model duplicates and synchronize the large amount of parameters.
16
PP and TP have been proposed to address the scalability issue. PP 72 16 partitions layers of a model across the GPUs. A batch is split into smaller microbatches, and training of different layers is pipelined with the microbatches across the GPUs. However, to ensure consistent weight updates for particular batch (without being affected by the weight updates from the other batches), GPUs need to synchronize the weight updates microbatches for every batch. This introduces pipeline bubbles where some GPUs (which process former microbatches of batch) need wait for the weight updates from the other GPUs (which process latter microbatches of the batch) delaying the entire training process 772. TP 4. on the other hand, distributes matrix multiplications within each transformer layer (i.e., MHA and MLP) across the GPUs. Each GPU handles portion of the matrix multiplications in parallel, without having model duplicates or pipeline bubbles, making highly effective for large-scale models.
Although TP receiving much attention recently for its scalability benefit further enhance memory efficiency and latency with large-scale models 5 55), its further efficiency is still limited by the communication overhead. Fig. 10 illustrates the train time and memory usage comparison between DP, PP and TP. While TP shows the fastest training time among three methods as it does not require communication of full parameters and pipeline bubbles, frequent communication between GPUs is still required to process synchronized and complete intermediate activations and gradients from MHA and MLP. As result, large portion of the training time devoted these communications (37.9% of the training time), resulting in a notable decrease in training efficiency.
To further enhance the potential of TP for fast large model training, propose FAL, which eliminates intra-block data communication by harnessing the output of the MHA in the first (i.e., bottom-most) transformer block for the MLP's inputs, instead of using the output of the MHA in the same block.
C Additional Motivation Analyses
C.1 Motivation Analyses in Different Scale
Fig. 11 shows the CKA similarity scores for MHA outputs, MLP inputs (Residual + MHA), and MLP outputs across adjacent blocks (conducted on GPT-2 774M and 1.5B). shown in Fig. even in the case of larger models, the MLP input remains highly similar despite significantly changing MHA
26

---

MHA oulpei
+U+ Residial
MHA ourgut - MLP oulpui
Wikitext-2
PTB
BookCorpas
CC-News
1.2
1.2
Simiarity
0.9
Simtlariy
0.9
0.6
0.6
0.3
0.3
12
24
35
16
Transformer Block Index
Transfommer Block Index
GPT2-L similarity analysis
GPT2-XL simi analysis
Figure 11: CKA-based similarity analysis of GPT-2 774M and 1.5B. across successive Transformer blocks. The x-axis shows the block index, and the y-axis shows the similarity (CKA) between consecutive MHA output, Residual MHA output (i.e., the MLP input), and the MLP input
output. This demonstrates that, regardless of the model size, MLP may not require the most recent MHA output (i.e., the output the MHA within the same block).
Original
ALL MBA
4800
4800
Perplesity
3600
3000
1200
Wikilext-2
PTB
BookCorpas
OC_ News
Wikitext-2
PTB
BookCorpus
CC News
(a) GPT2-L connection ablation
(b) GPT2-XL connection ablation
Figure 12: Connection ablation results measured by perplexity with GPT-2 774M and 1.5B, "Original" denotes the unaltered model. "All MHA" removes every MHA layer. "All Connect" removes every direct MHA-MLP connection.
illustrates two scenarios: removing all MHAs (All MHA) versus removing all MHA-MLP connections (All Connect), measured by perplexity on GPT-2 774M and 1.5B. As expected, removing All MHA severely degrades model quality. In contrast, removing All Connect recovers significant portion of the lost performance compared to removing the entire MHAs, and this recovery becomes pronounced with larger models (though still not fully reaching the original performance). This suggests that bypassing MHA is a better option for larger models.
Fig. 12
Wikiten-22
PTB
BookCorpas
OC-News
Gradient Magnitale
1.2
Gradient Migniiule
1.2
Normalized
Normalizd
0.8
0.8
0.4
0.4
16
20
24
36
28
36
44
48
Transformer Block Index
Transformal Block Index
(Gradient Magnitude MHA Outputi
(Gimdient Magninude MHA Outpur)
(3) GPT2-1 gradient analysis
GPT2-XL, gmddzet analysis
Figure 13: Normalized gradient magnitude of the MHA outputs across Transformer blocks in GPT-2
774M and 1.5B for different datasets. the x-axis represents the block index.
Fig. 13 shows the gradient magnitude of each MHA output larger scale (774M and I.5B). As shown in Fig. 13 even in the case of larger models, first MHA output consistently exhibits the highest gradient magnitude. This confirms our finding that perturbations in the earliest attention result have disproportionately large impact on final predictions, regardless of model size.
Fig. 14 shows the perplexity after omitting the MHA from individual transformer blocks with 774M and 1.5 scale GPT-2. As shown in Fig. [4 removing the first attention causes larger perplexity increase than removing later layers, verifying the crucial role of the first attention in language modeling. These findings align with the well-known psychological phenomenon the primacy effect [26), commonly summarized as "first impressions matter"" The primacy effect of the first
27

---

Wikitext-2
PTB BookCorpas
CC-News
1.3
1.3
Normaizal Perpicsity
1.2
Norna
Peppokity
1.2
1.1
16
20
36
48
Trimsoomer Block Index
lodex (where MHA layer is omittod)
(Where MHAA layer is omilted)
GPTZ-L layer ablation
(b) GPTZ-XL layer oblition
Figure 14: Layer ablation results measured by perplexity with GPT-2 774M and 1.5B, the x-axis indicates the index of the transformer block from which the MHA is omitted.
attention is not limited to specific model architecture previous works also identified the prominent impact of the first attention layer across various attention mechanisms and tasks 22 28 29
C.2 Motivation Analyses with Different Task & Model Architecture
Beyond scaling analyses, we further validate our motivation across different tasks and model archi- tectures, including ViT-B (86.6M), LLaMA2-7B, CodeLLaMA-34p3 and LLaMA3.2-11B. 1-B-1.33.--BBissooon
C.2.1 VIT
ontput - Residual
output -
ontput
CIFAR-100
CIP
CIFAR-10
Oxford Pet
Food-101
Original
ALL MHA
All Connect
1.1
100%
Similariyy
0.9
Acamcy
75%
0.7
50%
0.5
25%
0.3
(1%)
pet
Tannsfomer Block Index
VTT-B similarity analysis
ViT -B connection ablation
Figure 15: (a) CKA-based similarity analysis of ViT. across successive Transformer blocks. The x-axis shows the block index, and the y-axis shows the similarity (CKA) between consecutive MHA output, Residual MHA output (i.e., the MLP input), and the MLP input (b) Connection ablation results measured by accuracy. "Original" denotes the unaltered model. "All MHA" removes every MHA layer. "All Connect" removes every direct MHA-MLP connection.
Fig. 15 (a) shows the CKA similarity scores for MHA outputs, MLP inputs (Residual + MHA), and MLP outputs across adjacent blocks (conducted on ViT-B). As shown in the Figure, even in the case of vision task, the MLP imput remains highly similar despite the MHA output changing significantly. This confirms our findings that MLP may not always require the most recent MHA output (i.e., the output of the MHA within the same block), regardless of the domain.
Fig. 15 (b) illustrates two scenarios: removing all MHAs (All MHA) versus removing all MHA-MLP connections (All Connect), measured by accuracy VTT-B. As expected, removing All MHA severely degrades model quality. In contrast, removing All Connect recovers a large portion of the lost performance compared removing the entire MHAs, however this recovery becomes smaller with smaller vision models. This suggests that simply bypassing MHA on small-scale vision models may harm their accuracy.
Fig. 16 (a) shows the normalized gradient magnitude each MHA output in the encoder-base ViT-B 73 Although the effect is less pronounced than in language models, the first MHA output still consistently exhibits the highest gradient magnitude.
Fig. 16 (b) shows the accuracy after omitting the MHA from individual transformer blocks with ViT-B. As shown in the Figure, removing the first attention causes a far larger accuracy drop than removing later layers (except the last), verifying the crucial role of the first attention in language modeling. The
28

---

CIFAR-100
CIFAR-LO
Oxtord Pot
Food-101
Gradien: Magnituce
1.2
1.1
Narmalizad
0.8
Normalzed
Acsuracy
0.9
0.7
0.4
0.5
0.3
12
7
Transoruse Block Index
Trnnforner Black Index (where MIIA layer conitted)
(Grundient Magnitude MHA Output)
(4) ViT-B gradient analysis
(b) layer
Figure 16: (a) Normalized gradient magniude of the MHA outputs across Transformer blocks with ViT for different datasets. The x-axis represents the block index. (b) Layer ablation results measured by accuracy with ViT. The x-axis indicates the index of the transformer block from which the MHA is omitted.
importance of the final attention likely stems from the fact that VIT uses domain- classifier based on the output of the final Transformer block [73].
C.2.2 LLaMA
To further validate the generality of findings across modalities and scales, extend motiva- tional analyses to larger models: LLaMA2-7B (language), LaM-A.L (code generation), and 1.AMA3.2-11B-iistie (multilingual rision-language We use WikiText for LLaMA2-7B, The Stack
dataset 74 for CodeLLLMMA-34B, and the COCO captioning dataset for LLaMA3.2-1 1B-11BB- Vision.
Table 3: Similarity Analysis Results (Metric: CKA Similarity std)
Activation
LLaMA2-7B
CodeLLaMA-334 LLaMA32211B-Vision
Atin Out
0.60 10.14
0,66 #0.13
0.84 10.10
MLP In
0.98 #0.03
0.99 #0.03
0.98 #0.02
MIP Out
0.70 #0.23
0.80 10.18
0.80 #0.20
T the similarity analysis results, measured by CKA similarity (mean + std). Across all three models, we observe that the MLP inputs remain highly similar to each other (LLaMA2-7B: 0.98, CodeLLaMA-34BB 0.99, LLaMA3.2-1 1AA.2-11B--V 1B-VViion: 0.98), whereas the attention/MLP outputs vary more significantly (0.60-0.84) 048.. This analysis shows that the residual path already accumulates sufficient attention signals, making the MLP input less sensitive to the most recent MHA output. The same trend holds even for larger-scale (7B33BB) and multimodal (vision-language) models, extending earlier motivation analyses 3.1 and supporting the validity of reconfiguring MHA-MLP connections in FAL.
Table Layer Ablation vs. Connection Ablation (Metrie: Validation Perplexity)
Setting
LLaMA2-7B CodeL.LaMA-34B
LLaMA32-11B--istte
Original
7.39
1.80
25.12
Remove Layer
5339.99 99
1484.52
3995,84 84
Remove Connection
892.12
37.06
504.72
reports validation perplexity under two ablations: (1) removing the entire MHA layer, and (2) removing only the direct MHA-MLP connection. Across all three models, removing entire layers leads to catastrophic degradation (e.g.. PPL >> 1000), while removing only the connections recovers substantial portion of performance. Although the connection remmoedd models still fall short of the original baseline, they consistently perform far better than the layer- removed ones, and the recovery effect becomes more pronounced at larger scales. These results show that reconfiguring MHA-MLP connections causes far smaller degradation than removing entire layers, yet still falls short of fully matching the original performance. stable alternative signal needed, which is exactly what FAL. provides by reusing the impactful first attention.
29

---

Table Gradient Analysis (Metrie: Gradient L1 Norm)
Block
LLaMA2-7B
CodeLLaMA-33B
1st
505,99
321.65
2-End avg tstd
85.90
#102.00
46.24
+60.39
Ratio (1st/avg)
5.9x
7.0x
Table 5 reports the gradient analysis of MHA outputs in LLaMA2-7B and CodeL.aaM--34B, The first attention block shows much larger gradient magnitude (505.99 and 321.655 (655.))) compared to the average of later blocks (85.90 and 46.24, respectively). On average, the first block gradients are 5.9x and 7.0x larger than subsequent ones. These results highlight that the first attention exerts disproportionately strong influence on final predictions. Thus, reusing this impactful signal FAL remains well-justified across larger models and diverse tasks.
Table 6: Layer Ablation per Block (Metric: Validation Perplexity)
Block
LLaMA2-7B
30d4LLaMA-34B
LLaMA32-11B-
Original
7.39
1.80
25.12
1st
34.37
4.56
40.94
2-End avg tstd
4.37
+1.37
1.81 +0.01
24.75
#1.25
Ratio (1st/avg)
7.9x
2.5x
1.7x
Table reports the effect of ablating individual attention layers LLaMA2-7B, CodeI.LaMA-344B, and LLaMA3.2-10 IB-Visonn. Removing the first attention block causes a far larger degradation in validation perplexity (e.g.. 34.37 vs. 7.39 for LLaMA2-7B, 4.56 vs. 1.80 for CodeLLaMA-34e and 40.94 vs. 25.12 for LLaMA3.2-1 IB- (B-iiooo)) compared to ablating later layers, whose impact remains relatively small. The impact of removing the first attention is consistently larger than that later layers, up 7.9x in LLaMA2-7B. These results show that the first attention disproportionately important across scales and modalities. Its removal uniquely destabilizes the model, whereas later attentions contribute far less. This further supports that reusing the first attention in FAL remains well-justified across larger and more diverse models.
D Additional Evaluation Results and Analyses
D.1
Ablation study
We conduct two ablations to verify the effectiveness of reusing the first attention output in FAL and FAL+. Table 7 shows the validation perplexity and training time.
Table Comparison of validation perplexity and training time using Openwebtext dataset (GPT-2
774M)
Model
Perplexity
Training time
GPT-2 774M (Baseline)
17.75
13.2 days
FAL
17.55
8.6 days
FAL+
17.24
13.2 days
Ablation
21.34
13.2 days
Ablation2
17.98
8.6 days
Ablation1 (leveraging latest attention). Simple addition of normalized outputs within block degrades the quality, calling for the necessity of the first attention signal. We apply the same LN + LN structure from FAL using the latest attention output instead of the first attention.
X, + (LN(())(X)) + MLP,
+ LN(MHA,LN(XX
(3)
30

---

This leads to significantly higher perplexity (21.34) compared to the baseline (17.75), suggesting that the latest attention output does not provide stable or beneficial signal for MLP imputs under this reconfiguration.
Ablation2 (removing connection without first). Retaining only the first MHA-MLP connection is not sufficient to retain the first attention signal. We remove all MHA-MLP connections except for the first one.
XI + MHA,, + MLP, (LN((1) + MHA, (LN(XL)) [X++ (LN((()(()) AMP,(LL(XX,)
if
(4)
otherwise.
Although perplexity (17.98) remains comparable to the baseline, it still degrades model quality compared to FAL. This implies that merely connecting the first attention once is not sufficient to maintain the overall performance.
Comparison to FAL and FAL+. FAL reuses the first MHA output in each block. This reuse, aided by an additional LN to balance the first and residual signals, improves perplexity (17.55) and reduces training time (8.6 days). FAL+ augments the original connections with the first attention, achieving an even lower perplexity (17.24) but training time similar GPT-2. These results confirm that properly integrating the first attention output is crucial for both efficiency and model quality.
GPT-2.
X; + MHA((LN(X)) + MLN,(LN(X, + (HA,(LN(X))))
(5)
FAL.
X; + (LN(()(()))) + (LN((L)(X)) + (LN(HH,(LN(X)))
(6)
FAL+
X1 + (
+MLP, (LN(X,) + (LN()))
if ffi=1.
X,-1=
X; MHA, (LN(X))
+MLP,
(LL(P(LL(X.++ (LN(X;+ (LN(.)) LYNMHA(
otherwise.
Ablation with Other Layers. To further verify the benefit of reusing the impactful first attention, we compare FAL+ with variants that reuse the output of other MHA layers (2nd, 3rd, and so on). We adopt FAL+ in a 48-block configuration (see Fig. 9) and train for 500k steps under the same one-cycle schedule with a batch-size ramp-up to 8,192, ingesting 1.02B
Layer (FAL+)
Others
MIL Loss
tokens regardless of hardware speed.
3.5
shows the MLM loss comparison of FAL+, which reuses the first attention, against variants that reuse later-layer attentions. As shown in Fig. [7 reusing later-layer attentions consistently underperforms compared to the first attention. This implies that reusing weaker or less dominant signals is not as effective as leveraging the impactful first attention, confirming that the first attention provides uniquely stable and beneficial feature for reconfiguration.
Fig. 17
2.5
200M
400M 600M 800M IB
# of Ingested Tokens
Figure 17: Loss comparison using other MHA's output
D.2 How FAL Mitigates Information Dilution in Deep Transformers
In a standard Pre-LN Transformer, each block processes only its immediate predecessor's 's output, analogous to unrolling an RNN over depth. As depth grows, early-layer signals must traverse many transformations and risk dilution or loss-much like long RNN sequences forget initial states. FAL breaks pure sequential dependence by feeding the first-layer attention output directly into every later block, akin eelf-attention's ability to attend globally. Concretely, we normalize the first MHA output once and then add it to each block's 's input: ML smpptt = LN(X)) + LN (MA (LN (X+)))
31

---

FAL
FAL+
Relative weight
1.0
1.0
(for (forthee | MHHA
0.6
0.6
0.2
0.2
0.0
0.0
7
10
13
16
19
22 25 28 31 34 37 40 43 46
7 10
13
16
19 22
25
28 31 34 37 40 43 46
Transformer Block Index
Transformer Block Index
Figure 18: Relative LN scaling parameters connected to the first MHA output, normalized to average LN scaling across layers.
Residual connections alone cannot prevent the first signal from fading: activation variance ac- cumulates layer by layer [9], eventually washing out early-layer information much like how self-attention's (rr²) interactions dilute key dependencies over very long sequences (75.. FAL sidesteps this pitfall by reusing normalized first-attention tensor at each block, reinforcing the most salient initial context without extra overhead - echoing cognitive insights that rethinking first impression can lead to better decisions in deep reasoning [30].
Figure 18 shows the average LN scaling parameters (r) for the terms connected to the first MHA output, normalized by the average LN scaling across layers. Across depth, both FAL and FAL+ consistently assign non-negligible weights (roughly equivalent a 0.58:1-1:1 ratio compared the current block input), indicating that later blocks actively and adaptively incorporate the first- attentttnn signal after training. This dynamic weighting suggests that FAL can alleviate information dilution by adaptively reinforcing the initial signal across depth.
D.3 Inference Acceleration of FAL
GPT2
FAL
Computation
Communication
Normalized Time
Nornalized Time
0.8
0.6
0.2
of GPC
GfG
GPUS
GP
Sequence Length 1024
Sequence Length
2048
Sequense Length
1024
Sequence Length 2048
GPT-2 774M
(b) GPT-2
Normlizedd Time
Nzzd ired Time
0.88
0.88
0.6
0.6
GFG
GPUs
GPUs
GF
Sequence Length
1024
Sequence Length 2048
Sequence Length
1024
Sequence Length 2048
(c) GPT 2.5B
(d) GPT 8.3B
Figure 19: Normalized Multi-GPU Inference Time GPT-2 and FAL.
TP is often used to accelerate the inference execution, as it can accelerate per-request latency unlike other two parallelism methods (i.e., Data Parallelism and Pipelined Parallelism) 76). In order to evaluate the inference acceleration that can be further achieved by FAL in TP, we measure the forward step time without gradient calculation which aligned with the Time To First Token (TTFT) in the inference execution.
32

---

Fig. 19 shows the normalized forward step time of GPT-2 and FAL on multi-GPU server with NVLink (System 4), across model sizes ranging from 774M to 8.3B and sequence lengths of 1024 and 2048. As shown in the figure, tensor parallelism (TP) reduces per-request inference time by utilizing multiple GPUs. With 8 GPUs, TP achieves an average speedup of 56.5% for a sequence length of 1024 (up to 67. 67.6%). and 62.8% for sequence length of 2048 (up to 70.5%). However, as the number of GPUs increases, the degree of acceleration is limited by the communication overhead between GPUs. In such cases, FAL improves inference performance over GPT-2 by (1) significantly reducing inter-GPU communication and (2) increasing intra-GPU parallelism. Across configurations from to 8 GPUs, FAL reduces inference time by 11.1% average (up to 31.6%).
E Evaluation of Generalizability to Transformer Variants
E.1
Loss comparison using variants of Multi Head Attention
Pre-LN
FAL
FAL+
6
4.5
MLM Loss
MLM Loss
4
3.5
200M
400M
6ROM
S0OM
1B
200M
400M
600M
800M
1B
of Ingested Tokens
of Ingested Tokens
(n) FAL with GQA
(b) FAL with MoE-Attention
Figure 20: Loss comparison across different attention mechanisms: Grouped Query Attention (GQA),
and MoE-based Attention Moo-AAtention).
FAL and FAL+ can be applied to various Pre-LN based transformer variants, such as LLaMa (with Grouped Query Attention (GQA)), and Switch Transformer (with Mixture of Experts (MoE)) to improve the efficiency and model quality. To evaluate the generalizability of FAL and FAL+ to these variants, we measure token ingestion efficiency before and after applying FAL and FAL+ to GQA-based 777 and MoE-based attention models [58]. We adopt FAL and FAL+ in a 48-block configuration (see Fig. 9). To ensure consistent scheduling across hardware setups, we train for 500,000 steps using step-based one-cycle learning rate scheduler. The batch size is ramped up to a maximum of 8,192 by step 300,000, resulting in a total of 1.02B tokens ingested regardless of hardware speed.
Fig. 20 (a) shows the comparison of FAL and FAL+ when applied to GQA. Each attention layer uses GQA with two groups. This setup differs from standard Multi Head Attention (MHA) primarily in the key/value projections, and becomes equivalent to MHA when the number groups equals the number of heads. The results closely resemble those observed with standard MHA. The loss gap between the proposed architectures and the baseline remains consistent, demonstrating that the gains from FAL and FAL+ extend robustly to this efficient attention variant.
Fig. 20 (b) shows the comparison of FAL and FAL+ when applied to MoE-based Attention. We follow the configuration of MoE- based Attention introduced in the Switch Transformer, which was found to be unstable and thus not included the final architecture. Each expert in the MoE attention has its own query projection and tied key/value projections. One of two experts is activated per attention layer. Unlike the instability observed when using Switch layers attention, FAL and FAL+ do not suffer from gradient instability. The loss gap between the proposed architectures and the baseline remains consistent, demonstrating that the gains from FAL. and FAL+ extend robustly to sparsely activated MoE attention mechanisms.
33

---

Table 8: Comparison of validation accuracy using ImageNet dataset (ViT-B 86.6M)
Dataset
ViT (Baseline)
FAL
FAL+
ImageNet
79.06%
78.76%
79.20%
E.2 Accuracy comparison using other kinds of Task
We also evaluate FAL and FAL+ on the Vision Transformer (ViT-B 86.6M) architecture using the ImageNet dataset. As shown in Table 8 FAL slightly reduces accuracy compared to the baseline (79.06% vs. 78.76%) whereas FAL+ achieves higher accuracy (79.20%) This result indicates that reusing the first attention output can also benefit vision tasks, particularly when combined with the original attention connections as in FAL+.
34


Reviews:
[{"review_id": "I8aNrZonUO", "reviewer": "Reviewer_eLP5", "summary": "This paper proposes **FAL** (First Attentions Last), an architecture that removes each block’s\nMHA-MLP link and instead feeds the first attention output, which has been normalized, to all\nlater MLPs. This halves per-block all-reduce traffic in Tensor Parallelism, since the all-reduce\ncommunication between each MHA-MLP has been replaced. This enables parallel MHA and\nMLP execution and up to 44\\% faster multi-GPU training (\\(\\le\\!1.18\\times\\) single-GPU\nthroughput). A variant **FAL+** augments rather than replaces the link, yielding even lower\nperplexity. Overall, the approach cuts communication, improves efficiency, and slightly\nimproves or preserves language-model quality versus GPT-2 baselines.", "questions": "- The first attention tensor is reused in every subsequent block; how much additional activation\nmemory does this introduce for long sequences and deep models? In the paper, minimizing\noverhead (including memory usage) is mentioned as one advantage, but it would be better if it\nis tested.\n- Has this model been tested with an even-larger scale LLM-models or vision-language\nmodels, in which case the \"first impression matters\" or \"similarity stability\" may fail? It is good\nthat in motivation, similarity analysis and connection ablation hold for this case, but it would be\nbetter if the generality to other cases are tested.\n- During instruction tuning or RLHF, where later layers often adapt more than early ones, does\nrepeatedly injecting a frozen first-layer attention harm adaptation or cause catastrophic\nforgetting? It would be better if it is tested, and consider whether dynamic first-layer attention\nwill work better", "limitations": "Here is one suggestion to improve the insight of the paper. The idea of this paper is impressive:\nutilizing the first attention to later MHA-MLP connections to get the intermediate activation can\ngreatly reduce overhead, including synchronization waiting and memory usage. However, the\nreason of the potential improvement should also be conjectured, or even proved. Intuitively, we\ncan interpret the improvement by regarding original Pre-LN transformer as a recurrent neural\nnetwork, where the output from the previous layer is passed to the next one in a sequence.\nHowever, FAL, proposed in this paper, is similar to attention mechanism, which imitates the\ntransformer to relate MHA-MLP attention with first attention feature (not relate to other\nattention layers due to \"first impression matters\" in the paper). Therefore, we may conjecture\nthat: the reason why FAL outperforms traditional methods in some cases may follow a similar\nreason why transformer generally outperforms recurrent neural networks in. If possible, this\nidea can be evaluated through experiments or reasoning.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n- *Reasonable motivation.* After realizing single MHA output change does not influence the\nMLP input greatly, and the output from the first attention layer dominates instead, the need of\nall-reduce in each block may be unnecessary.\n- *Elegant insight.* Re-uses the first‐layer attention as a cheap and effective substitute for\nevery MHA-MLP link.\n- *Practical speedups.* Reports up to 44\\% end-to-end training time reduction on 4×PCIe\nGPUs and $1.18\\times$ single-GPU throughput, with the same PyTorch code path.\n- *Quality preserved.* Perplexity drops slightly and zero-shot SuperGLUE scores hold or rise,\nespecially with the FAL$^{+}$ variant.\n\n\n**Weaknesses**\n- *Limited task breadth.* Experiments centre on GPT-2 style LM pre-training and SuperGLUE;\nvision or multilingual models are untested.\n- *Statistical rigor.* All metrics come from single runs, and the variance of the performances are\nnot tested.\n- *Memory analysis absent.* Holding the first attention tensor through $N$ blocks may raise\nmemory needed to cache the features, but this is not measured nor discussed.\n- *Appendix deficiency.* Some contents, like analyzing the limitations, are not included in the\nmain contents, which would be better to be mentioned in some parts.\n- *Downstream robustness.* No fine-tune tasks (instruction tuning, RLHF, retrieval‐augmented\ninference) are benchmarked, so practical quality impact is under-explored.\n- *Model scale ceiling.* Experiments stop at $8.3$ B parameters and $\\le 60$ layers; effects on\ntruly large LLMs ($\\ge 50$ B, hundreds of layers) remain untested.", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your valuable feedback.\n\n### **W1&Q2:** Limited Task Breadth\nTo further demonstrate its cross-model applicability, we provide the following evidence:\n\n\n- For a **vision model (ViT)**, as in Appendix C.2 and E.2, we observe that the similarity of MLP inputs (0.9409 on average) remains consistently higher than that of MHA outputs (0.7859) or MLP outputs (0.7637). Additionally, the gradient magnitude of the first attention layer is 5.42× larger than that of other layers. Leveraging this signal, FAL+ achieves a +1.4% accuracy gain over the ViT baseline.\n\n\n- We have also newly extended our motivational analyses to larger-scale models, including **CodeLLaMA-34B (LLM)** and **VisionLLaMA-11B (vision-language, multilingual)**, further supporting the generality of our findings across modalities and scale. **The results shown in the below tables align with our motivation analyses conducted on GPT-2 and ViT. This 1) demonstrates high similarity in MLP inputs and 2) highlighting the pivotal impact of the first attention on overall model quality**. We will incorporate these new findings into the paper.\n\n\n\n>Similarity Analysis Results (CKA Similarity ± std)\n>\n\n| Activation | LLaMA2-7B | CodeLLaMA-34B | VisionLLaMA-11B |\n|------------|----------|---------------|-----------------|\n| **Attn Out** | 0.60 ±0.14 | 0.66 ±0.13 | 0.84 ±0.10 |\n| **MLP In**   | 0.98 ±0.03 | 0.99 ±0.03 | 0.98 ±0.02 |\n| **MLP Out**  | 0.70 ±0.23 | 0.80 ±0.18 | 0.80 ±0.20 |\n\n>Layer Ablation vs Connection Ablation (Validation PPL)\n>\n\n| Setting | LLaMA2-7B | CodeLLaMA-34B | VisionLLaMA-11B |\n|---------|----------|---------------|-----------------|\n| **Original** | 7.39 | 1.80 | 25.12 |\n| **Remove Layer** | 5339.99 | 1484.52 | 3995.84 |\n| **Remove Conn.** | 892.12 | 37.06 | 504.72 |\n\n>Gradient Analysis (Gradient L1 Norm)\n>\n\n| Block | LLaMA2-7B | CodeLLaMA-34B |\n|-------|---------:|--------------:|\n| **1st** | 505.99 | 321.65 |\n| **2~End avg ±std** | 85.90 ±102.00 | 46.24 ±60.39 |\n| **Ratio (1st/avg)** | 5.9× | 7.0× |\n>Layer Ablation per Block (Validation PPL)\n>\n| Block            | LLaMA2-7B | CodeLLaMA-34B | VisionLLaMA-11B |\n|------------------|---------:|--------------:|----------------:|\n| **1st**          | 34.37 | 4.56 | 40.94 |\n| **2~End avg ±std** | 4.37 ±1.37 | 1.81 ±0.01 | 24.75 ±1.25 |\n| **Ratio (1st/avg)** | 7.9× | 2.5× | 1.7× |\n\n\n- Note, even for the model variants, such as GQA and MoE-based attention, FAL and FAL+ maintain consistent loss improvements over the baseline (see Appendix E.1).\n\n\n\n### **W2:** Statistical Rigor \n**Performance Measurements:** All reported speedups are averaged **over 50 runs** after an initial warm-up period to stabilize thermal and caching effects, on four representative platforms.  Across these settings, the **relative deviation of measured training/inference times is below 0.01%**, confirming that our reported gains are robust and not due to measurement noise. We will include error bars in the final paper to reflect the variance across multiple runs.\n\n\n**Model Quality:** Although we have not conducted exhaustive seed swaps due to the limited GPU resources available, we believe the model quality results will be consistent --- **pretraining and zero-shot performance evaluation across three distinct corpora---OpenWebText, SuperGLUE, and Pile—exhibit consistent perplexity and accuracy patterns in every case**.\n\nIf we get access to more GPU resources in near future, we will conduct the experiments with more random seeds to further ensure the consistent conclusions --- we have tried to get the access during the rebuttal period, but we failed due to contentions.\n\n\n\n### **W3&Q1:** Memory Analysis Absent.\nFAL does not incur any additional memory overhead during training --- it rather slightly reduces total memory usage compared to the Pre-LN Transformers.\n\n\n**No Extra Activation Buffers:** To prevent storing additional activations—which would be required if the original MHA output were transformed (e.g., via multiplication, addition, or compression)—we **reposition** the first layer’s Layer Normalization (from MLP’s input to MHA’s output). This enables later blocks to reuse the output of the first layer's Layer Normalization without any recomputation as well as memory overhead. Note, every MHA output (including the first layer’s) is already cached during the forward pass to enable gradient computation.\n\n\n**Fewer LayerNorms → Slightly Reduced Memory:** FAL **reduces** total memory usage by approximately 1% compared to the baseline, primarily by reducing the number of LayerNorm operations and associated intermediate activations. Instead of applying a fresh LayerNorm in every block, FAL normalizes the first-attention output only once (in block 1) and reuses the result across all blocks.\n\n\n\n### **W4:** Appendix Deficiency\nWe will add the concise summarization of the Appendix contents (such as key limitations) in the main text.\n\n\n\n### **W5&Q3:** Downstream Robustness\nAs the reviewer suggested, we empirically examined how injecting the first attention signal during instruction tuning affects the trade-off between adaptation and stability — note, given limited rebuttal period (and resources), we have not evaluated RLHF setting in this rebuttal due to its high memory and time requirements. We fine-tune both GPT-2 1.5B and FAL+ 1.5B on the **Alpaca instruction tuning dataset** using four different learning rates (1e-5 to 1e-2). We report:\n\n\n**Trained perplexity on Alpaca**, reflecting adaptation, and\n\n\n**Δ Validation perplexity on OpenWebText**, reflecting forgetting.\n\n| Model        | LR   | Δ Val PPL | Trained PPL |\n|--------------|-----|---------:|-----------:|\n| **FAL+‐1.5B** | 1e-5 | **0.00** | **27.97** |\n|              | 1e-4 | **0.00** | **27.98** |\n|              | 1e-3 | **-0.01** | **27.22** |\n|              | 1e-2 | **0.61** | 5.76 |\n| **GPT2‐1.5B** | 1e-5 | 0.01 | 30.93 |\n|              | 1e-4 | 0.00 | 30.92 |\n|              | 1e-3 | 0.01 | 28.10 |\n|              | 1e-2 | 1.39 | **4.83** |\n\n\n**FAL+ consistently preserves pretraining knowledge better than GPT-2.** Across all learning rates, FAL+ shows lower or equal ΔPPL on OpenWebText compared to GPT-2, indicating stronger stability. At LR=1e-3, FAL+ even shows a slight improvement in validation PPL (-0.01), suggesting positive regularization rather than forgetting.\n**FAL+ also achieves better adaptation when forgetting is minimal.**  While maintaining zero or negative ΔPPL, FAL+ attains the lowest trained PPL (27.22) at LR=1e-3, demonstrating high adaptability with minimal forgetting. In contrast, GPT-2 reaches a lower trained PPL (4.83) at LR=1e-2 only by severely compromising original domain knowledge (ΔPPL = 1.39).\nWe agree with the reviewer that **dynamically injecting the first attention** (e.g., through gating or reweighting) is an interesting direction. However, such mechanisms would require **storing additional intermediate activations** and introducing **block-specific parameters**, resulting in **non-trivial memory and compute overhead**. Balancing this added complexity against adaptability gains would require careful exploration, which we leave for future work.\n\n\n\n### **W6:** Model Scale Ceiling\nWhile our main experiments are done on a relatively small scale compared to SOTA LLMs, both **efficiency gains and quality preservation/improvement are scalable to larger models**. For the details, please see our response to **W2&Limitations** of **Reviewer R6Kr**. Note we have conducted experiments with 34B CodeLLaMA instead of a 50B model due to the limited resources. Our request (to run the experiments of 50B model on 8 H100 GPUs) is waiting in the queue of the cloud server. If we finish executing the experiments, we will update the results in the paper.\n\n### **Response to Suggestion:** How FAL Mitigates Information Dilution in Deep Transformers\nWe thank the reviewer for this insightful framing and let us answer your suggestion.\n\n**Sequential “RNN‐like” Information Dilution:** In a standard Pre-LN Transformer, each block processes only its immediate predecessor’s output, analogous to unrolling an RNN over depth. As depth grows, early‐layer signals must traverse many transformations and risk dilution or loss—much like long RNN sequences forget initial states.\n\n\n**FAL as a Persistent “First-Impression” Skip:** FAL breaks pure sequential dependence by feeding the first-layer attention output directly into every later block, akin to self-attention’s ability to attend globally. Concretely, we normalize the first MHA output once and then add it to each block’s input:\nMLP_input_i = LN(X_i) + LN(MHA_1(LN(X_1)))\n\n\n**Preventing Signal Dilution in Deep Transformers:** Residual connections alone cannot prevent the first signal from fading: activation variance accumulates layer by layer (Xiong, Ruibin, et al., 2020), eventually washing out early-layer information —much like how self-attention’s O(n²) interactions dilute key dependencies over very long sequences (Yao, Shunyu, et al., 2021).  FAL sidesteps this pitfall by reusing normalized first-attention tensor at each block, reinforcing the most salient initial context without extra overhead—echoing cognitive insights that rethinking a first impression can lead to better decisions in deep reasoning (Wason et al., 1974).\n\n - Xiong, Ruibin, et al. \"On layer normalization in the transformer architecture.\" International conference on machine learning. PMLR, 2020.\n- Yao, Shunyu, et al. \"Self-attention networks can process bounded hierarchical languages.\" arXiv preprint arXiv:2105.11115 (2021).\n- Wason, Peter C., and J. St BT Evans. \"Dual processes in reasoning?.\" Cognition 3.2 (1974): 141-154.\n\n**Empirical Weighting Ratios:** Analysis of the learned LN‐affine parameters shows that blocks weigh the first-attention signal at ratios ranging roughly from 1 : 1 to 1 : 2 compared to the current (block) input, confirming its active, dynamic role across depth. We will include detailed per-layer ratio plots in Appendix D."}, {"review_id": "rrIy9ksvAM", "reviewer": "Reviewer_EX1R", "summary": "To alleviate the communication overhead in TP training, this paper proposes FAL, a transformer architecture that selectively skips MHA for saving allreduce communications in TP, and FAL+, a variant that sends the first MHA output to intermediate layers through residual connections. The evaluations shows training speed improvement due to saved communication on GPT-2 model.", "questions": "1. Fig. 3(b) : was it trained from scratch, or fined tuned on a trained model?\n2. If the first MHA output is important, how about layer #2, #3, etc, and combinations of them? On top of Figure 4, it may be worthwhile to provide an ablation study for FAL that uses more than just layer 1.", "limitations": "Yes", "rating": 3, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strength:\n* This paper identifies a important problem with tensor paralleism -- frequent collective communication -- and proposes a novel solution that selectively skips intermediate layers. \n* The evaluation on GPT-2 shows promising results on performance improvement, especially on cases where slow interconnect dominate training.\n* The paper is well-written and easy to follow. \n\nWeakness:\n* The key motivation for the MHA-skipping design is based on GPT-2 only, and it is not clearly justified whether the observation that \"MHA-MLP connections can be selectively skipped\" stands on other models like LLama or Qwen. \n* Some claims are flawed. For instance, line 256 states FlashAttention increases computational intensity and gives FAL more oppurtunity for overlapping compute-bound and memory-bound ops. This is not very precise as MLP is compute bound too, and this actually reduces the chance for overlapping. And FAL does not actually *overlap* them unless it can fuse the the (skipped) parallel transformer layers.", "was_revised": true, "final_justification": "The authors have address my concerns and I would like to maintain my score.", "rebuttal": "We appreciate your valuable feedback.\n\n### **W1:** Motivation Generality Beyond GPT-2\nOur core motivation—similarity after residual addition and the dominant role of the first attention—applies broadly to **Pre-LN-based Transformer architectures**, which underlie most modern large-scale models.\n**Vision Transformer (ViT-B)**: In Appendix C.2, we analyze the normalized gradient magnitudes on ViT-B and observe that the first MHA output remains the main contributor and the MLP input remains highly similar despite the MHA output changing significantly even in the case 518 of vision task, mirroring our GPT-2 findings.\n**Attention Variants (GQA & MoE)**: Appendix E.1 reports fast-training loss comparisons for Grouped Query Attention and MoE-based attention in a 48-block setup; FAL and FAL+ maintain consistent loss improvements over the baseline, confirming adaptability to these modern variants.\n\nTo further validate cross-model applicability, we have extended our **motivation analysis to**:\n\n\n**Llama2-7B** (7 B parameters)\n\n\n**CodeLlama-34B** (34 B parameters, Code Generation)\n\n\n**VisionLlama-11B** (11 B parameters, multilingual vision-language)\n\n>Similarity Analysis Results (CKA Similarity ± std)\n>\n\n| Activation | LLaMA2-7B | CodeLLaMA-34B | VisionLLaMA-11B |\n|------------|----------|---------------|-----------------|\n| **Attn Out** | 0.60 ±0.14 | 0.66 ±0.13 | 0.84 ±0.10 |\n| **MLP In**   | 0.98 ±0.03 | 0.99 ±0.03 | 0.98 ±0.02 |\n| **MLP Out**  | 0.70 ±0.23 | 0.80 ±0.18 | 0.80 ±0.20 |\n\n>Layer Ablation vs Connection Ablation (Validation PPL)\n>\n\n| Setting | LLaMA2-7B | CodeLLaMA-34B | VisionLLaMA-11B |\n|---------|----------|---------------|-----------------|\n| **Original** | 7.39 | 1.80 | 25.12 |\n| **Remove Layer** | 5339.99 | 1484.52 | 3995.84 |\n| **Remove Conn.** | 892.12 | 37.06 | 504.72 |\n\n>Gradient Analysis (Gradient L1 Norm)\n>\n\n| Block | LLaMA2-7B | CodeLLaMA-34B |\n|-------|---------:|--------------:|\n| **1st** | 505.99 | 321.65 |\n| **2~End avg ±std** | 85.90 ±102.00 | 46.24 ±60.39 |\n| **Ratio (1st/avg)** | 5.9× | 7.0× |\n\n\n\n\n>Layer Ablation per Block (Validation PPL)\n>\n\n| Block            | LLaMA2-7B | CodeLLaMA-34B | VisionLLaMA-11B |\n|------------------|---------:|--------------:|----------------:|\n| **1st**          | 34.37 | 4.56 | 40.94 |\n| **2~End avg ±std** | 4.37 ±1.37 | 1.81 ±0.01 | 24.75 ±1.25 |\n| **Ratio (1st/avg)** | 7.9× | 2.5× | 1.7× |\n\n\n\nThese experiments confirm (1) the feasibility of bypassing per-block MHA outputs, and (2) the pivotal role of the first MHA output across architectures and domains.\n\n\n### **W2:** Overlap of Compute‐and Memory‐Bound Operations\nFAL executes the MHA and MLP modules on separate CUDA streams, allowing their operations to proceed concurrently. This allows the warp scheduler to hide latency effectively—when one warp stalls (e.g., on memory), another ready warp from a different stream can be scheduled without delay. We further clarify the interaction between MLP and attention phases as follow:\n\n**1. MLP GEMM still incurs memory stalls.**\n\nAlthough MLP’s large matrix multiplies are compute-heavy, each GEMM begins and ends with global‐memory loads and stores (Global Memory → Shared Memory → Registers → Cuda/Tensor Cores → Shared Memory → Cuda cores → Global Memory). These boundary memory transactions introduce unavoidable stalls, even if the core GEMM is compute-bound.\n\n\n**2. FlashAttention creates “slack” for warp‐scheduling.**\n\nFlashAttention’s higher arithmetic intensity in the attention block lengthens its compute phase relative to its memory phase. Thus, by extending attention’s compute window, FlashAttention generates more opportunities to hide MLP’s boundary stalls behind attention computation.\n\n**3. Empirical GPU‐utilization gains.**\n\nFigure 8-(b) shows that FAL with FlashAttention achieves higher sustained utilization of SMs and memory pipelines compared to the Pre-LN baseline, confirming that attention’s heavier compute helps to hide MLP memory stalls.\n\n\n**4. Stronger effect in tensor-parallel settings.**\n\nIn multi-GPU tensor-parallel training, MLP GEMMs are split across ranks, shortening the compute window of MLP. FAL’s attention-compute overlap mechanism becomes even more effective here, since the added memory latency in TP is better masked by FlashAttention’s extended compute phase.\n\n\nThese points will be clarified in Section 6.3 of the revised manuscript.\n\n\n### **Q1:** Fig. 3(b) Model Status\nAll analyses in Figure 3(a) and 3(b) use an **off-the-shelf, pre-trained GPT-2** checkpoint—**no fine-tuning** was performed. Specifically, we load the HuggingFace GPT-2 (774 M) checkpoint and run it on WikiText-2, PTB, BookCorpus, and CC-News to study its internal behavior in a stabilized setting. As noted in Appendix A.1, no additional training or task-specific fine-tuning occurs for these experiments.\n\n\n### **Q2:** Fig. 3(b) Ablation with Other Layers\nWe thank the reviewer for the suggestion. When we reuse the attention of later layers instead of and along with the first one, we observe the perplexity degradation even compared to the GPT-2 baseline.\n\n**Latest‐Layer Variant (Ablation 1 in Appendix D.1):** we replace the first‐layer signal with each block’s own MHA output. This degrades validation perplexity to 21.34 versus 17.75 for the GPT-2 baseline, confirming that the first attention is uniquely informative.\n\n**Layer‐2&3:** Under fast‐training (“Cramming”) conditions at 48 layers (like Figure 9 - Middle), using the second and third‐layer output atop the first one yields an MLM loss of 2.39, 2.58, respectively, worse than the Pre-LN baseline (2.32) and both FAL (2.22) and FAL+ (2.17). This implies that **additionally injecting less dominant signals along with the first one adversely affects understanding of input during training**.\n\nWe will include these ablation study results in Appendix D.1.\n\nNote, in addition to the perplexity degradation, injecting later-layer signals also necessitates additional intra-block all-reduce operations --- this is directly counter to FAL's design goal to halve the communication overhead."}, {"review_id": "u6jwZhamkA", "reviewer": "Reviewer_R6Kr", "summary": "This paper addresses the computation overhead in large language model pretraining, particularly within the context of tensor parallelism (TP). The authors noticed that the all-reduce operation in the connection of multi-head attention (MHA) and feed-forward layer (MLP) forms a bottleneck. Based on two empirical findings: 1) MHA-MLP connections can be reconfigured and 2) first-layer attention output is more important, the authors propose two new architectures FAL and FAL+, which add further skip-connection for each MLP layer with the first-layer attention outputs. Experiments with GPT-2 show that the proposed method is more compute efficient and could achieve better performance than the baseline.", "questions": "1. In large-scale modeling, the residual connection is unlikely to be a major bottleneck, as major computation is spent on the attention and MLP layers. It would be significant to show timing analysis with model scaling.\n2. As FAL/FAL+ requires storing the first MHA output, they may incur further memory overhead. Could you provide a detailed analysis from the memory perspective? This could be significant as sequence length, batch size and model scale increases.\n3. In Table 1, at 1.5B, FAL+ achieves better PPL but worse SuperGLUE. Could you provide some potential reasons for this?", "limitations": "The observation, like 44% training time reduction, is likely not generalizable to real larger models.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths:\n\n- the analysis is intuitive and the method is easy to understand and implement\n- FAL and FAL+ achieve decent performance on several evals, either improved efficiency and/or better quality, compared to the baseline\n\n\nWeaknesses:\n- GPT-2 is a relatively old architecture. Findings may not generalize well to modern architectures, like LLaMa and Gemma.\n- Model scale is generally too small, under which the conclusion is unlikely to generalize to larger setups.", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your valuable feedback.\n\n### **W1:** Generality Beyond GPT-2\nTo demonstrate FAL’s broad applicability, we include the following additional evaluations:\n\n**Transformer Variants (GQA & MoE): Appendix E.1** shows that applying FAL/FAL+ to Grouped Query Attention —as used in LLaMA and Gemma— and to MoE-based attention yields loss-reduction curves and stable quality gains nearly identical to those on standard MHA, confirming that our “first-attention matters” insight extends to these modern variants. **Appendix E.2** reports ImageNet accuracy for ViT-B (86 M) where FAL+ actually **improves** top-1 accuracy (79.20% vs. 79.06%) and FAL matches baseline, indicating cross-domain robustness.\n\n**Visaon-Language and Code Models:** To further validate cross-model applicability, we have extended our **motivation analysis to:**\n**Llama2-7B** (7 B parameters)\n**CodeLlama-34B** (34 B parameters)\n**VisionLlama-11B** (11 B parameters, multilingual vision-language)\n\nThe experimental results confirm (1) the feasibility of bypassing per-block MHA outputs, and (2) the pivotal role of the first MHA output across architectures and domains.\nFor the detailed results, please see our response to **W1&Q2** of Reviewer **eLP5**.  \n\n**Supporting Literature:** Prior works such as (Liu, Liyuan, et al., 2020, Liu, Zichang, et al., 2023) report minimal activation changes across layers, and (Chen, et al., 2024, Zhang, et al., 2024) the important role of first MHA which are consistent with our findings. These extensions firmly establish that FAL’s principles—exploit the first MHA output to reduce overhead by replacing subsequent direct MHA–MLP connections— hold across modern architectures, like LLaMa and Gemma.\n\n- Liu, Liyuan, et al. \"Understanding the difficulty of training transformers.\" arXiv preprint arXiv:2004.08249 (2020).\n- Liu, Zichang, et al. \"Deja vu: Contextual sparsity for efficient llms at inference time.\" International Conference on Machine Learning. PMLR, 2023.\n- Chen, Xingwu, Lei Zhao, and Difan Zou. \"How transformers utilize multi-head attention in in-context learning? a case study on sparse linear regression.\" Advances in Neural Information Processing Systems 37 (2024): 119573-119613.\n- Zhang, Yang, Yanfei Dong, and Kenji Kawaguchi. \"Investigating layer importance in large language models.\" arXiv preprint arXiv:2409.14381 (2024).\n\n### **W2&Limitations:** Model Scale\n\nWhile our main experiments are done on a relatively small scale compared to SOTA LLMs, both **efficiency gains and quality preservation/improvement are scalable to larger models**.\n\n**1. Performance Scalability — Model-Independent Gains from Communication Reduction**\n\n\nBy cutting per-block communication in half, FAL can achieve on the order of **44 % (or greater) end-to-end training time reduction** irrespective of model size, and its gains are expected to become more pronounced for larger scales.\n\nFAL’s key benefit comes from halving the frequent all-reduce operations between each block’s MHA and MLP — precisely the communication bottleneck in Tensor Parallel training. Such communication overhead is expected to be more significant in large-scale setups, as communication bandwidth has not kept pace with compute performance; over the last two decades, raw FLOP/s per chip grew by roughly **60 000×**, whereas interconnect bandwidth improved by only **30×** (Gholami et al., 2024). This growing gap already manifests in practice. For instance, **training a 200 B-parameter model on TPU v4 Pods** shows that **communication alone accounts for 42% of training time** (Wang et al., 2022). On typical GPU-based systems — where interconnect bandwidth is **up to 10× slower** than TPU v4 (Selvan & Kanwar, 2021) — this fraction is expected to be even higher. FAL’s architecture is designed precisely for such scenarios, making it especially valuable in **large-scale GPU clusters** with slower interconnects, where communication overhead dominates. As such, FAL’s performance improvements remain robust across scales and are likely to **amplify on future systems** as the compute/communication imbalance continues to widen.\n\n**2. Model Quality — Scaling trends show greater benefit at larger depths:**\n\n\nQuality benefits of FAL and FAL+ become more pronounced for large-scale, deeper models.\n\n\n**Deep-Model Scaling Experiments:** When we increase the model depth in fast-training scaling tests (36 → 48 → 60 layers as in Fig. 9), both FAL and FAL+ consistently reduce masked-language-model (MLM) loss more rapidly than the Pre-LN baseline. While all models converge to similar losses at 36 layers, FAL and FAL+ achieve progressively lower loss as the model depth increases.\n\n\n**Rationale behind the Deep-Model Scaling Results:** In standard Pre-LN architectures, activation variance accumulates across layers, gradually washing out early-layer information (Xiong, Ruibin, et al., 2020) — **residual connections alone cannot fully preserve** the initial signal. FAL and FAL+ counteract this attenuation by repeatedly reintroducing the high-impact first MHA output at every block, **allowing each block to adaptively re‑weight and reinterpret this signal** throughout training.\n\n**Key Motivation Generality to Large-Scale Models:** When we extended our motivation analysis to larger models, **VisionLlama-11B** and **CodeLlama-34B**, we observed the same two key phenomena: **feasibility of bypassing** per-block MHA outputs without catastrophic quality loss and **pivotal role of the first MHA output** as the critical bridging signal for performance recovery. For the detailed analysis results, please see our response to **W1**.\n\n**Together**, as model depth—and thus the “distance” between input and final output—increases, repeatedly leveraging the first attention signal yields ever greater quality gains. This makes FAL and FAL+ especially valuable for state-of-the-art, large-scale transformer models.\n\n- Gholami, Amir, et al. \"Ai and memory wall.\" IEEE Micro 44.3 (2024): 33-39.\n- Wang, Shibo, et al. \"Overlap communication with dependent computation via decomposition in large deep learning models.\" Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1. 2022.\n- Selvan, Aarush, and Pankaj Kanwar. \"Google showcases Cloud TPU v4 Pods for large model training.\" Google Cloud Blog (2021).\n- Xiong, Ruibin, et al. \"On layer normalization in the transformer architecture.\" International conference on machine learning. PMLR, 2020.\n### **Q1:** Timing Analysis & Bottleneck Attribution\n\nAlthough the majority of GPU time is devoted to attention and MLP compute, multi-GPU data-movement—the all-reduce between each block’s MHA and MLP—still accounts for **25.4–31.7%** of total training time across scales (GPT-2 774M: 31.79%; 1.5B: 25.40%; 2.5B: 31.35%; 8.3B: 31.65%, on H200/NVLink). \n\nFigures 6 and 7 break down end-to-end time into:\n**Compute** (MHA + MLP forward/backward)\n**Communication** (all-reduce)\n\nFAL’s connection reconfiguration targets **multi-GPU data-movement** rather than residual-path compute. By halving the number of per-block all-reduce calls, FAL directly cuts the communication fraction, yielding up to **44% end-to-end training time reduction** (Fig. 6).\n\n\n### **Q2:** Memory Overhead of Reusing First Attention\n\n\nFAL does not incur any additional memory overhead during training --- it rather slightly reduces total memory usage compared to the Pre-LN Transformers.\n\n\n**No Extra Activation Buffers:** To prevent storing additional activations—which would be required if the original MHA output were transformed (e.g., via multiplication, addition, or compression)—we **reposition** the first layer’s Layer Normalization (from MLP’s input to MHA’s output). This enables later blocks to reuse the output of the first layer's Layer Normalization without any recomputation as well as memory overhead. Note, every MHA output (including the first layer’s) is already cached during the forward pass to enable gradient computation.\n\n\n**Fewer LayerNorms → Slightly Reduced Memory:** FAL **reduces** total memory usage by approximately 1% compared to the baseline, primarily by reducing the number of LayerNorm operations and associated intermediate activations. Instead of applying a fresh LayerNorm in every block, FAL normalizes the first-attention output only once (in block 1) and reuses the result across all blocks.\n\n\nUnlike FAL, which is designed to minimize the communication cost without incurring additional memory overhead, **FAL+ is designed primarily to enhance model quality** by injecting additional first attention signals in each block.This augmentation introduces a small memory overhead of approximately\n\n\n((1 × B × L × H) for activation + (2 × H) for affine parameters) × (#Transformer blocks − 1)\n\n\nwhere B is the batch size, L the sequence length, and H the hidden dimension — resulting in about a 1% increase in memory usage.\n\n### **Q3:** Discrepancy Between PPL and SuperGLUE at 1.5B\nWe attribute this discrepancy to a domain mismatch between the pretraining corpus (OpenWebText) and the downstream SuperGLUE benchmark. While perplexity measures language modeling quality on the pretraining distribution, SuperGLUE evaluates generalization to a diverse set of reasoning and classification tasks. Prior work (Tay, Yi, et al., 2021) has shown that improvements in perplexity do not always translate to downstream gains in the presence of domain gaps. \n\nNevertheless, FAL+ achieves better SuperGLUE scores than GPT-2 at both 774M and 1.5B, despite slightly underperforming FAL in the 1.5B case. One possible explanation is that FAL's more constrained signal path acts as a form of implicit regularization.\nWe believe that continued pretraining on more diverse and downstream-aligned corpora can further close this performance gap.\n\n- Gururangan, Suchin, et al. \"Don't stop pretraining: Adapt language models to domains and tasks.\" arXiv preprint arXiv:2004.10964 (2020)."}]

Author rebuttal:

Meta-review: This paper proposes FAL, which greatly saves multi-GPU communication overhead by removing MHA-MLP communication bottlenecks. Reviewers appreciated the careful ablations and clear presentation - there was initially some concern about potential specificity/overfitting to the GPT-2 architecture, but the authors addressed this concern adequately in their rebuttal. Overall, the scientific rigor and practical results makes the paper of interest to the broader community.
Decision: Accept (poster)