Weaknesses

Quality

Applying this method to dense transformers further demonstrates that the gate enables stable training with larger batch sizes and learning rates, resulting in improved performance.` It is difficult to make this claim definitively...

Thank you for this suggestion. We increased the batch size when increasing the number of training tokens because we aimed to compare under the best possible baseline performance. Our hyperparameter search for the baseline revealed that both batch size and learning rate need to be increased as the number of training tokens grows, which is consistent with existing scaling laws (Predictable Scale: Part I, Step Law -- Optimal Hyperparameter Scaling Law in Large Language Model Pretraining). We understand that this setup may introduce additional confounding factors (i.e., token count) when comparing the effect of gating under different batch sizes. Therefore, we have further added experiments under setting (c):

Method	LR	Tokens	Bsz	Avg PPL	MMLU	HumanEval	GSM8k	Hellaswag	C-eval	CMMLU
Baseline	4e-3	400B	1024	7.499	50.21	28.66	27.82	64.94	49.15	49.52
SDPA Elementwise	4e-3	400B	1024	7.404	51.15	29.27	28.28	65.48	50.72	50.72

Baseline	4e-3	400B	2048	7.538	49.02	24.37	25.72	63.54	48.02	48.82
SDPA Elementwise	4e-3	400B	2048	7.487	50.41	27.56	27.13	64.11	49.64	50.03

We observe that increasing the batch size larger than the optimal one leads to a performance drop. However, adding the SDPA output gate still results in better performance than the baseline, with more pronounced improvements on benchmarks such as MMLU, HumanEval, and GSM8k.
Nevertheless, we acknowledge that the phrase resulting in improved performance could be misleading and that the comparison remains challenging. In the next version, we will revise this statement to also resulting in performance improvement relative to the baseline.

Significance

For example, overall, results in Table 1 are not significant performance improvements. It is difficult to say "significant reduction in PPL" in Table 3. Similarly, it is difficult to call a 0.2 PPL reduction a significant performance improvement...

We would like to provide additional context to demonstrate that a 0.2 PPL reduction is a meaningful performance gain:
In our Table 2, for the 48-layer model, increasing training tokens from 400B to 1T leads to only a 0.06 reduction in PPL for the baseline (from 7.4 to 7.34), while the SDPA Elementwise gate achieves a 0.187 reduction (from 7.28 to 7.093). In contrast, in Table 1, the gate reduces PPL from ~6.0 to ~5.8—a 0.2 drop—even though architectures differ, this remains a substantial improvement.
Additionally, we supplement Table 1’s setting with baseline models trained for 500B and 600B tokens (with full cosine schedulers), as shown below:

Method	Tokens	Avg PPL	MMLU	GSM8k
Baseline	400B	6.026	58.79	52.92
Baseline	500B	5.851	60.17	54.85
Baseline	600B	5.704	61.03	55.15

SDPA Elementwise	400B	5.761	60.82	55.27

Notably, the SDPA Elementwise model trained on 400B tokens outperforms the baseline trained on 500B tokens and is only slightly behind the 600B baseline.
For further reference, in the Fox paper (mentioned by the reviewer), Table 3 shows that adding an output gate in a 360M model trained on 7.5B tokens reduces PPL from 7.08 to 6.88—a 0.2 drop. While test PPLs are not directly comparable due to different datasets, this supports our claim that a 0.2 PPL reduction is meaningful.
Moreover, we will remove the word significant to avoid mis-understandings.

For dense models, while gating-placement was much more impactful on 48-layer 1T pretraining token models, gains are modest for the 28-layer model. It is possible that substantial improvements may have occurred for the missing experiments above.

Thank you for raising this point. We have further supplemented experiments on the 28-layer 1.7B model trained on 400B tokens with increased learning rates. The results are shown in the table below:

Method	LR	Avg PPL	HumanEval	MMLU	GSM8k	Hellaswag	C-eval	CMMLU
Baseline	4e-3	7.499	28.66	50.21	27.82	64.94	49.15	49.52
SDPA Elementwise	4e-3	7.404	29.27	51.15	28.28	65.48	50.72	50.72
Baseline	8e-3	8.941	20.13	41.17	16.15	58.10	44.36	45.13
SDPA Elementwise	8e-3	7.425	30.37	52.51	29.24	65.30	51.76	51.88

As shown, when the learning rate is further increased, the baseline model encounters convergence issues and suffers significant performance degradation. In contrast, the model with SDPA Elementwise gating maintains stable training and achieves further improvements across all benchmarks except PPL and Hellaswag. This supports our earlier claim that the gating mechanism enables stable and effective training under larger learning rates.

Originality

"Despite its widespread adoption and empirical success, the function and impact of gating mechanisms remain insufficiently explored beyond their initial intuition." This claim is too broad, the impact of gating has been explored in linear attention [1] and standard attention [2,3] networks.

Thank you for the feedback. We have discussed the relation to [2] in line 300. In the revised version, we will add the following clarifications regarding [1] and [3]:
[1] discusses the difficulty of Mamba’s decay mechanism in forgetting historical hidden states and also uses an output gate at the SSM layer. However, neither [1] nor the original Mamba2 paper analyzes the behavior or impact of the output gate in detail.
[3] proposes a Gumbel-Softmax-based gate in LSTM attention to mask irrelevant tokens for the current query. In contrast to [3], which applies a hard mask before SDPA to explicitly remove irrelevant token information, our attention output gate operates after SDPA and uses gating scores to suppress irrelevant information, effectively mitigating the attention sink phenomenon.
We will modify the expression to Despite its widespread adoption and empirical success, most recent works do not look into the gating mechanisms like the gating scores and their effect on the model’s hidden states.

Questions:

Why were the same placement experiments in Table 1 not repeated for the dense model in Table 2? The missing "Max LR" experiments for different configurations (e.g., (9) SDPA Headwise trained with max LR 8.0 × 10⁻³) make it difficult to draw definitive conclusions.

Initially, we focused on evaluating gating positions in MoE models due to their lower training cost and strong performance, expecting similar trends in dense models. We now supplement the results for different gating placements in a 28-layer, 1.7B dense model trained on 400B tokens. These will be added to the appendix:

Method	Avg PPL	MMLU	GSM8k	Hellaswag	C-eval	CMMLU
Baseline	7.499	50.21	27.82	64.94	49.15	49.52
SDPA Elementwise G1	7.404	51.15	28.28	65.48	50.72	50.72
v Elementwise G2	7.453	50.53	27.97	65.12	49.82	49.91
k Elementwise G3	7.485	50.25	28.01	64.99	49.03	49.17
q Elementwise G4	7.466	50.33	27.89	65.04	49.15	49.33
Dense Output G5	7.489	50.34	28.08	65.01	49.38	49.25

We further supplement results for SDPA Headwise at LR = 8e-3:

Method	Avg PPL	MMLU	GSM8k	Hellaswag	C-eval	CMMLU
SDPA Elementwise	7.325	54.47	36.62	66.40	53.91	53.80
SDPA Headwise	7.384	54.13	34.56	65.97	54.10	52.38

We observe that under more aggressive training setting, the elementwise gate outperforms the headwise variant, and the training curve is smoother. Hence, we recommend the elementwise gating design.

Incorrect citations

"Prior work has established that while increased network depth, large learning rates, and large batch sizes can significantly improve model performance [33, 34, 35] and distributed training efficiency, they often introduce training instabilities [34, 36, 37]."

Citations 34 and 36 do not support the above claim. 34 attributed instability to pre vs post residual placement of the norm, while 36 did not attribute instability to any particular factor (but rather explored layernorm placement's effect on stability).

Thank you for pointing this out. The confusion likely arose from combining multiple factors across citations.
For [34], we intended to use its Table 1 to support the claim that increased depth can cause instability (many models diverge when depth increases), and its Table 2 to support that depth can improve performance (both baseline and DeepNet show gains with more layers).
For [36], a more appropriate citation would be A Theory on Adam Instability in Large-Scale Machine Learning, which thoroughly discusses how larger learning rates and batch sizes can lead to loss spikes and other instabilities. [36] primarily proposes techniques like Embedding Layer Gradient Shrink (EGS) to address such spikes .
We will modify the expression accordingly.

Thanks again for your suggestions!