Masked Gated Linear Unit

We thank the reviewer for acknowledging the effectiveness of our MGLU layer. We respond to the reviewer’s questions below.

The latency comparison can be unfair, and speed-up can be overclaimed. The 18x speed-up numbers come from comparison to the pytorch implementation, which is very inefficient. Some commonly used techniques, like kernel fusion and CUDA Graph, can largely improve the baseline performance.

To ensure a fair comparison, we have now implemented the same kernel-fusion and CUDA-Graph optimizations enabled with torch.compile in our baseline. Tables below summarize the updated micro-benchmark results on a RTX5090 and a H100 GPU. Even against an optimized MGLU baseline, FlashMGLU still yields up to 12.23× speedup in the CUDA kernel.

RTX 5090

H100

What is the end-to-end model inference speed-up?

With a JIT-fused model torch.compile(..., mode="max-autotune", fullgraph=True), our FlashMGLU CUDA implementation enabled end-to-end decoding speed-up is around 1.12$\times$ on our 1B model on a RTX 5090.

We sincerely appreciate your valuable feedback.

What would be the overall runtime breakdown, and how much gain can we get?

The runtime of the original LLM consists of SwiGLU FFN (50.2%), self-attention (43.6%), and other operations (6.2%), indicating that the SwiGLU FFN is the computational bottleneck. Since our SwiMGLU accelerates the up-projection layer of the SwiGLU FFN by 1.5$\times$, the overall inference speed improves by 12% while also improving downstream performance.

It seems the performance still lags behind the SwiGLU, I may be more convinced when seeing larger models’ results.

Pre-training large-scale LLMs from scratch demands substantial GPU resources and trillion-scale token datasets. For instance, Llama 3.1 8B was trained using 24K H100 GPUs on 15 trillion tokens, totaling 39.3M GPU hours, which is not feasible within our research budget. Nevertheless, we have made our best efforts to train an 8B-parameter LLM on a dataset of 100B tokens (FineWeb-Edu) to address the reviewer’s concern (although the data is insufficient), utilizing 128 GPUs over three days. The results demonstrate that our SwiMGLU outperforms SwiGLU by a clear margin at this stage: 3.11 vs. 3.17 in training loss, 22.2 vs. 23.4 in perplexity, and 42.02 vs. 40.50 in downstream performance. This suggests that our approach is expected to scale effectively to larger models.

Hyperparameter (branch number) may be tricky to choose.

Based on the compatibility of current GPU architectures with our CUDA kernel implementation, we recommend selecting $m=4$. While the optimal number of branches depends on the downstream task, this behavior is consistent with observations reported in previous studies on MoE architectures. Given that using multiple branches ($m>1$) consistently improves performance, our claims regarding the efficiency and effectiveness of MGLU remain valid.

Limited validation of the proposed method, which is conducted on a small-scale model. More model series and larger model scales should be conducted.

Results using GeGLU (as in the Gemma series) are provided in the Appendix, demonstrating that our method remains consistently effective. Scalability is discussed above.

There seems to be no theoretical connection as to why sharing the same weights by a mask can obtain similar performance as SwiGLU.

Sharing the same weights by using a binary mask achieves similar performance to SwiGLU because, theoretically, a sufficiently large random weight matrix inherently contains multiple distinct, smaller "submatrices" capable of performing separate transformations. By applying a learned binary mask, the single weight matrix can be effectively partitioned into two complementary parts, each approximating the role of one of the original SwiGLU matrices. Recent theoretical results inspired by the Lottery Ticket Hypothesis (E. Malach, et al., ICML 2020 [1]) confirm that if the shared matrix is made just slightly wider, it becomes highly likely to contain such suitable submatrices that approximate the two independent matrices used in SwiGLU. Additionally, since the activation function (such as Swish) is smooth, small differences in these approximations lead only to minor differences in the output. Consequently, a masked weight-sharing layer can replicate the expressivity and accuracy of SwiGLU at a significantly lower memory footprint.

[1] E. Malach, et al. Shamir, Proving the Lottery Ticket Hypothesis: Pruning is All You Need, ICML, 2020.

We hope our responses have adequately addressed the reviewer's concerns.

We thank the reviewer for the insightful comments. We are encouraged that the reviewer found our idea smart and well-executed. Below we respond to each comment.

The authors forgot to include the appendices, so it is hard to properly judge the validity of the training and evaluation pipelines.

The appendix is attached separately as a zip file instead of being included in the main pdf (as permitted by the NeurIPS submission guidelines). To download it, please click the "zip" icon under the "Supplementary Material" section.

What are the wall-clock pretraining times of the experiments in section 5?

The wall-clock pretraining times of the experiments are described on Table 6 in Appendix A. Similar to other approaches like post-training quantization, we prioritized reducing inference costs, even though this leads to increased training time.

Do you have any hypotheses as to why MGLU leads to more robust and stable training than GLU?

As shown in our ablation (Appendix B.3) and in the mask distribution analysis (Appendix C), fixing the mask vectors in advance, instead of co-learning them with the weights, leads to substantially worse convergence and final accuracy. We therefore hypothesize that jointly learning the masks confers two key benefits.

1. MGLU removes high-variance activation fluctuations. This "smoothing" of the forward pass reduces gradient noise, allowing more consistent updates and preventing early-training oscillations.
2. The MGLU architecture inherently optimizes how resources are split between gate, value projections. Although the change in this allocation across layer depth is small, a clear trend emerges (see Figure 14), which may be the reason of MGLU yielding robust training.

You said that using inverse masks in motivated by the implementation efficiency, but have you tested if there aren't some downstream-performance losses compared to learning the two masks independently?

We've tested implementations using independent masks instead of complementary masks on large models. The table below illustrates the downstream performance on each model configuration for $n_m=4$. Our results show minor or no degradation in downstream performance due to using complementary masks.

The results suggests that the improvement in downstream performance get smaller with larger model scales. How do you expect your method to scale? Would MGLU still outperform GLU when the transformer is scaled to ~100B size?

While training a model in the ~100B parameter range is beyond the resources of our academic lab, we conducted an additional experiment on an 8B parameter model to further probe this scaling behavior. Given our computational constraints, we trained both SwiGLU and SwiMGLU variants on a 100B token subset of the training data, which was the maximum scale we could afford for this experiment, requiring nearly 10,000 GPU-hours.

The results, shown in the table below, reinforce our hypothesis and demonstrate the clear benefits of our approach. The SwiMGLU variant achieves a lower training loss, a better validation perplexity, and a higher average downstream score, indicating superior learning dynamics and generalization even on this larger scale.

Based on the consistent trend from our 1B model and these new 8B model results, we expect this architectural advantage to hold, and for MGLU to continue outperforming standard GLU variants as we scale to much larger models.

The paper should clearly state what is the computational cost of MGLU during training. I am happy to increase the rating after that is transparently communicated.

Thank you for your positive comment and advice. We will include the computational cost of training in the final version of the main text upon acceptance.

We appreciate the reviewer's careful reading and detailed feedback. Below, we respond to each question.

Line 163 states that the IO is $n_m \times 2 + 1$, but it becomes $n_m \times 4$. Could the author clarify this?

The description around Algorithm 1 mis-states the naive I/O count. The correct issued memory reads are $n_m \times 2 + 1$, where the full matrix vector multiplication $Wx$ (one matrix load) and the masked weight matrix vector multiplication $(W{\odot}M){\cdot}x$ (one weight and one mask matrix load per mask) are computed. The complimentarily masked matrix vector multiplication results are subtracted from $Wx$.

With $n_m = 8$, a fused kernel that reduces IO from 8 x 4 = 32 to 2 is supposed to save 16x. Could the authors explain how the actual speedup (>19x) is even higher than this estimated upper bound?

From the above discussion, the memory reads in $n_m = 8$ are (16bit x 9 + 8bit x 8) x numel = 208 bit x numel as boolean values are stored in int8 data types. This could be reduced to (16bit + 8bit) x numel = 24 bit x numel which should result in a 8.6x speedup. However in reality, the kernel launch time would be dominant especially in smaller models with multiple matrix operations, which could lead to even higher speedups.

I'm surprised that a CUDA implementation can be 6x faster than the Triton implementation. Could the authors elaborate on which aspects of the kernel Triton lacks? I would assume it's the lack of support for sub-byte tensors.

Our CUDA implementation uses explicit warp-level reductions and vectorized loads. In contrast, triton does not support sub-byte memory operations and memory loads are completely managed by triton runtime, which collectively account for about 3x gap on our implementation.

Line 254, what is the speedup of 1.51x compared against? (Is it standard SwiGLU or GLU?)

This 1.51x is measured against the standard GLU, on an RTX 5090 with $n_m = 1$, hidden = 8192, intermediate = 2048. We’ll add a parenthetical pointer in the caption for clarity.

Lines 265-266, am I correct in understanding that the end-to-end speedup compared to existing architecture ranges from 1.15x to 1.24x?

The 1.15x to 1.24x speedup range refers only to the isolated MGLU kernel microbenchmarks (Appendix E, Table 17, 18) for $n_m = 8$. We will revise the text to avoid confusion.

Line 194: Is it 3hd or 2hd? Line 195: Is it 2hd or 16hd?

We thank the reviewer for catching these mistakes. Line 194 should read 2hd as defined in Equation 3 and consistent with Table 2. Line 195 should read 16 hd. We will correct the final version of the paper.

Additional questions:

Are the masks input-dependent or independent (it seems to be the latter)?

The learned masks in MGLU are input-independent. Each mask parameter is learned jointly with the weight matrices, and does not vary with the input token.

Could the authors discuss the training stability when using the straight-through estimator for the masks? (I'm aware of such a section in 5.3, but I'm generally curious if using straight-through during pretraining is relatively straightforward.)

In our experiments, MGLU models converge at essentially the same rate as SwiGLU baselines. Additionally, we observe no spikes attributable to the STE. Norm and loss remain stable across layers even when we observe loss spikes on SwiGLU baselines.

Could the authors clarify what the main benefit of MGLU compared to SwiGLU is? Is it primarily for memory savings (and speedup during to reduced IOs)?

While GLU provides strong expressivity via input-dependent gating, MGLU’s primary advantage is in memory-bandwidth reduction and speedup during memory-bound inference while maintaining downstream performance.