Title: First Attentions Last: Better Exploiting First Attentions for Efficient Parallel Training
Venue: NeurIPS 2025
Accepted: yes
Keywords: Neural Networks

Parallelism

Distributed Training

Tensor Parallel
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/


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe observation, like 44% training time reduction, is likely not generalizable to real larger models.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-04"}

comments: 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.

STRENGTHS AND WEAKNESSES:
Strength:
* This paper identifies a important problem with tensor paralleism -- frequent collective communication -- and proposes a novel solution that selectively skips intermediate layers. 
* The evaluation on GPT-2 shows promising results on performance improvement, especially on cases where slow interconnect dominate training.
* The paper is well-written and easy to follow. 

Weakness:
* 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. 
* 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.

QUESTIONS:
1. Fig. 3(b) : was it trained from scratch, or fined tuned on a trained model?
2. 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
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\n- 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\n\nLIMITATIONS:\nHere 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.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-06"}