Masked Gated Linear Unit
Authors: Yukito Tajima
Nakamasa Inoue
Yusuke Sekikawa
Ikuro Sato
Rio Yokota
Venue: NeurIPS 2025 poster ()
Keywords: Activation
Sparsity
GPUs
LLM training
LLM inference
Primary area: deep_learning

Abstract: Gated Linear Units (GLUs) have become essential components in the feed-forward networks of state-of-the-art Large Language Models (LLMs).
However, they require twice as many memory reads compared to feed-forward layers without gating, due to the use of separate weight matrices for the gate and value streams.
To address this bottleneck, we introduce Masked Gated Linear Units (MGLUs),  a novel family of GLUs with an efficient kernel implementation.
The core contribution of MGLUs include:
(1) the Mixture of Element-wise Gating (MoEG) architecture that learns multiple binary masks, each determining gate or value assignments at the element level on a single shared weight matrix resulting in reduced memory transfer, and (2) FlashMGLU, a hardware-friendly kernel that yields up to a 19.7$\times$ inference-time speed-up over a na\"ive PyTorch MGLU and is 47\% more memory-efficient and 34\% faster than standard GLUs despite added architectural complexity on an RTX5090 GPU.
In LLM experiments, the Swish-activated variant SwiMGLU preserves its memory advantages while matching—or even surpassing—the downstream accuracy of the SwiGLU baseline.

Masked Gated Linear Unit
Yukito Tajima'
Nakamasa Inoue'
Yusuke Sekikawa²
Ikuro Sato 1.2
Rio Yokota'
1 Institute of Science Tokyo, Japan
Denso IT Laboratory, Japan
yukito@rio scrc. iir ac.
Abstract
Gated Linear Units (GLUs) have become essential components in the feed-forward networks of state-of-the-art Large Language Models (LL.Ms). However, they quire twice as many memory reads compared to feed-forward layers without gating, due the use of separate weight matrices for the gate and value streams. To address this bottleneck, we introduce Masked Gated Linear Units (MGLUs), novel family GLUs with an efficient kernel implementation. The core contribution of MGLUs include: (1) the Mixture of Element-wist Gating (MoEG) architecture that learns multiple binary masks, each determining gate or value assignments at the element level single shared weight matrix resulting in reduced memory transfer, and (2) FlashMGLL MU,,, a hardware-friendle kernel that yields up to 19.7x inference-time speed-up over naïve PyToorch MGLU and is 47% more memory-efficient and 34% faster than standard GLUs despite added architectural complexity on an RTX5090 GPU. In LLM experiments, the Swish-activated variant SwiMGLU preserves its memory advantages while matching-or even surpassing- downstream accuracy of the SwiGLU baseline.
Introduction
Transformers (Vaswani et al. 2017) have revolutionized deep learning and given rise to large- scale language models (LLMs) that achieve remarkable results across a wide spectrum of natural
language processing tasks (Brown et al. 2020
Wei et al., [022] Ouyang et al. 2022] Grattafioni
ct al. 2024] OpenAl| [0224] Gemma Team 20255 Yet the explosive growth in parameter count and
inference scaling translates directly into substantial inference and latency. The decode latency is fundamentally dominated by communication between high-bandwidth memory (HBM) and SRAM due transferring model weights for computation, therefore memory reduction being crucial in real-world applications (Gholami et al. 2024).
Modern LLMs predominantly adopt decoder-only architecture (Radford al. 2018). with each decoder layer typically consisting two modules: a self-attention module and a feed-forward network (FFN), Recent successful optimizations for attention mechanisms, such as FlashAttention (Dao et al.)
2022
Dao 2024 Shah et al., (024) and alternative approaches like the Mamba (Gu and Dao
2025
Dao and Gu 2024), have significantly improved the computational efficiency by reducing glboal memory reads/writes. Nevertheless, kernel-level optimization of the FFN remains particularly challenging, as its simple architecture provides fewer opportunities for efficiency improvements compared to attention.
The activation function within the FFN plays crucial role in determining the decoder output. Traditional nonlinearities such as ReLU (Narr and Hinton 2010) and GELU (Hendrycks and Gimpel 2016) laid the groundwork for deep learning, but the Gated Linear Unit (GLU) variants (Dauphin et al. 2017a] |Sazzerr 2020) introduce multiplicative interactions that markedly boost expressivity. For example, SwiGLU, which swaps the sigmoid gate in the original GLU for the smoother Swish function (Ramachandran et al.) 2017). has become a primary choice in modern LLMs because
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Learnable
Learnable
Weight
W
W
W
W
Mask M
W
GELU
Swish
Swish
M
W
W3
W
M&W
MeW
MoW
(a) GELU
(b) SWiGLU
(c) SwIMGLU (Single-Mask)
Figure 1: Comparison of FFNs. (a) Two-layer FFN using GELU. (b) SwiGLU FFN with gating mechanism that requires two separate weights: W, and W,- (c) Single-mask variant of SwiMGLU FFN (ours), which introduces a learnable binary mask M to decompose a single weight W into two complementary projections, reducing the required memory load during inference.
further enhances convergence and downstream accuracy. However, the explicit gating operations involving two separate projections and their and associated memory overhead inference time can erode throughput gains, especially latency-sensitive scenarios.
In this paper, we propose Masked Gated Linear Units (MGLUs). novel family of GLUs, along with an efficient CUDA kernel implementation. Rather than maintaining two separate projections Wg. We for the gate and value streams, our MGLUs adopt the Mixture of Element- wiss Gating (MoEG) architecture, which applies learnable binary masks M, to a single projection W, reproducing GLU's hallmark gating interactions. In Figure [1 we illustrate single-mask variant, where single weight W performs both gate and value projections via a mask M.
The core mechanism enabling efficient kernel implementation, which we call FlashMGLU, lies in using complementary masks M, = M4. Given an input vector æ, FlashMGLU simultaneously computes the gate a M, @ W) and the value ar M; @ W) by leveraging their complementarity With four masks and Swish activation applied, our SwiMGLU variant achieves downstream performance comparable to or even better than the widely adopted SwiGLU, while reducing computational cost of the projection layers by 29. 1 1% and memory usage by 37.5% during inference. It is worth noting that FlashMGLU is 12.51x x faster than a naive Py implementation of MGLU, opening new opportunities for more efficient and effective implementations of FFNs.
Our primary contributions are summarized follows:
We propose MGLUs, a novel family of GLUs with the MoEG architecture. We effectively mimic and enhance gating mechanism without incurring the computational and memory overhead associated with two separate full-rank projections.
We introduce FlashMGLU, efficient CUDA kemel implementation of MGLUs, reducing memory bandwidth requirements and enabling faster inference with minimal code modification.
variety of downstream NLP tasks, demonstrating that SwiMGLU achieves comparable or superior downstream accuracy to SwiGLU while notably improving inference throughput and memory efficiency, validating its practical effectiveness for resource-constrained LLM deployments.
We conduct extensive experiments
Related Work
Efficeint Inference via Sparse Masks. Model pruning methods seek to reduce the inference memory load by climinating redundant weights. Early unstructured pruning methods removes individual parameters achieve high sparsity (Frankle and Carbin 2019 Xia et al. 2022). While unstructured pruning yields irregular patterns that impede efficient execution, structured pruning excises entire components for straightforward speedups the cost of coarse- pears--rriined weight removal (Hou et al.) 2025; [andri et al.] 2025| Ashkboos et al. 2024). Semi-structured approaches blend these two extremes by enforcing regular N:M sparsity blocks, combining fine-grained flexibility with hardware- friendly patterns (Fang et al. 2024 Sun et al.) 2024 Frantar and Alistarh 2023). Learnable N:M domain specific specialized masks optimized on calibration data to attain up to 50% sparsity with minimal perplexity degradation and marked improvements in inference throughput. However, the

---

weights that are masked out still occupy storage but take no part in computation, leaving reservoir of unused capacity. If we can reactivate or repurpose these dormant parameters controlled way, we may further boost accuracy without increasing the deployed model's memory footprint.
Efficeint Inference via Activation Sparsity. Whereas weight sparsity reduces model size, activation sparsity directly lowers runtime FLOPs at inference (Zhang et al., 2025 HAZIZA et al.. 2025
Mirzadeh et al. 2024). Deja Vu shows that, for each input, only small, input-depender subset of heads and MLP channels is required; a lightweight predictor selects them on-the-fly, halving latency without gre-training or quality loss (Liu al., 2023). Complementary to this contextual approach, TEAL applies magnitude- based pruning to the hidden states themselves, achieving 40-50% uniform sparsity without any retraining and realizing decoding speed-ups on modern architectures (Liu et al., 2025). Activation sparsity methods reduce computation and memory load by skipping unneeded activations during inference. In contrast we compress weight matrices directly by adding arithmetic complexity in order to generate multiple streams of outputs.
Choices of Activation Functions in Large Language Models. Transformer (aswwani al. 2017)
MLP/FFN layers initially adopted simple piecewise-linear rectifiers such as ReLU (Nair and Hinton 2010), and Gaussian Error Linear Units (GELUs) (Hendrycks and Gimpel 2016) which weight inputs by the Gaussian CDF to produce smooth, non-saturating transitions. Later, non-monotonic self-gating activations like Swish (1-sigmoid(z))) were discovered via automated search, offering improved deep network performance through enhanced gradient propagation (Ramachandran et al. (017)) More expressive gated linear units (GLUs) introduce multiplicative interactions by splitting a projection into value and gate streams modulated via sigmoid (Dauphin et al. 2017b; Shazeer] 2020); variants such as GEGLU (using GELU) and SwiGLU (using Swish) have been incorporated into Llama series (Grattafiori et al. 2024) to boost convergence and downstream accuracy (Shazeer 2020).
3
Method
This section presents the Masked Gated Linear Units (MGLUs), novel family of GLUs that reduces memory access by emulating the gate and value streams using single shared weight matrix. We first review the GLU variants and discuss the challenges associated with sharing their two separate projections. We then introduce MGLUs, which replace explicit gating with a mixture of element-wise gating (MoEG), sculpting distinct subspaces from the single weight matrix through learnable binary masks. Although our MoEG increases the architectural complexity, it is designed to enable efficient CUDA kernel implementation, eliminating extra matrix multiplies and memory accesses.
3.1 Preliminary
GLU Variants. The GLU variants for FFNs (Dauphin et al. 2017a Shazeer 2020 augment a standard two-layer FFN by splitting the intermediate projection into the gate and value streams. Specifically, given an input vector æ E g² and two learnable weight matrices Wg, Wo E hxd the generalized GLU laye is defined as:
(1) where g is gating function, @ denotes Hadamard product, h is the hidden size, and d is the intermediate size. The resulting intermediate representation is then mapped back to the hidden size by an output projection W. E gd Figures ((a) and ( illussrrate the architectures of the standard
GLU(x) g(* o (=W.);
two-layer Linear Unit (LU) FFN using GELU (Hendrycks and Gimpel. 2016) and the SwiGLU FFN using Swish (Ramachandras et al. 2017) for g. respectively.
Can two matrices be shared in SwiGLU? While SwiGLU is a primary choice in state-of-the-are L.MSs (Grattafiori et al. 2024), the separate value and gate projection matrices incur 2x the memory reads compared to single linear layer. One might ask whether a single shared weight W can serve both streams, for example by using distinct channel-wise transformations or learned offsets. However, naively sharing W typically col- lapses the expressivity of the multiplicative interaction: independent full-rank projections are required
FFN
Parameters PPL
GELU
W
25.6
SwiGLU
{ 9. W W. 27.0
W. 23.6
SwiGLU (shared)
Table 1: Perplexity comparison.
'Trooughout this paper, we refer to the generalized GLU layer simply as the GLU layer unless ambiguity arises. Following recent practice in LLMs, omit bias vectors, but incorporating them is straightforward.
3

---

Learnable
Gate
Value
W.
M;
Moo
M,s
M2
MjoW
MyoW
Swish
Swish
W
MOW
MOW
MOW
M,O
M2
McoW
McaW
Mr Mco
McoW
Multi-mask SwiMGLU
(a)
Figure 2: Mixture Wiement-wise Gating (MoEG). All gate and value projection matrices are computed from the shared weight matrix W. (b) The MoEG-based SwiMGI.U architecture with m routes, each of which leverages element-wise gating-
to learn disentangled feature gating and value transformations shown in Table ! This could be limitation of explicit gating, motivating us to propose MGLUs with element-wise gating.
3.2
MGLU Layer
Our goal retain the expressive multiplicative effect of the GLU layer while reducing the number of its full-rank projections. Instead of the two separate projection matrices (W,, W. our MGLU layer introduces small set of binary masks M (M, 1=) each of which sculpts different subspaces of single weight matrix W.
Single-Mask Variant. We start by describing the single-mask variant of the MGLU layer. Let W E R be shared weight matrix and M {0, 1 be a binary mask matrix. We define the MGLU layer as
MGLU, g0 @ W)) @ (a(M @ W)).
where æ is an input vector, is a gating function, and M = M is the complementary mask. As shown in Figure [(c] the MGLU layer recovers the multiplicative gate-value interaction using only one full-rank weight matrix W. Through clement-wise gating with two binary matrices {M, M, W effectively partitioned into two complementary subspaces, resulting in greater parameter efficiency compared to the GLU layer. During training, M is optimized jointly with W via the straight-through estimator (Bengio et al. 2013) on the binarization.
We utilize the complementary mask rather than two separate masks because it enables more efficient CUDA kernel implementation. FlashMGLU in Section computes the gate projection æ [M @ W) and the value projection a I @ W) simultaneously by leveraging their complementarity.
MoEG Variant. To capture diverse gating patterns across channels, the MoEG architecture allows multiple complementary masks and lets the model learn which subspace W serves as gate versus value. Specifically, we define the MoEG-based MGLU as follows:
Am
MGLU n... (x)
Σ
g(x(MM W)) (a(M.
where nm. is the number of mixtures, and M Tm is set of binary mask matrices.
Figure shows the overall architecture. As shown, all gate projection matrices M, O and value projection matrices M, O W are derived from the shared full-rank weight matrix W (Figure and their resulting representations are subsequently aggregated (Figure (b).. This architecture, featuring m parallel routes, can be interpreted as variant of the mixture of expert architecture (Jacobs et al. 1991. Jordan and Jacobs 1994), thereby improving representational capacity. During training, the masks {M, 15 are optimized jointly with W. inference, all masks are fixed, and the fused masked projections execute with single kernel.

---

Input
On-chip SRAM
load load
Weight W
write
Sum
O
W
h
load
81
M,
0
Sub
gate streams
82
M4
Mask .M
M1
M
write
split, 2
Repister
772 value streams
Figure 3: Diagram of how FlashMGLU forward pass is performed during generative inference. When the weight W is split into two blocks in the K dimension (split, 2), the input vector is also partitioned into two. We pre-compute the unmasked matrix-vectes operation, than selectively add-up the according to the mask values avoiding excessive memory reads of weight matrices.
Algorithm FlashMGLU forward pass: Split-K K Matrix Vector Product with Packed Nym Masks.
Require: A R MMN, RR mask (0..... .0...., 2mm 13 MXN HBM, split,
I: Initialize accumulators E R X. f-
2: for row 0 to M do
split, 6- << + split, 1) //pppp 1. start 4- chunk x 8, end <- min(start s,8 N) 5: Each thread sets 1 < 0 and <- 0 for 1,...,
3: for chunk
6:
for each start.
end - 1 (stride -llocDDmm.r) do Load A[row, AJ, xx]], mask[row, k] from HBM register.
7:
On chip, compute
A(row, A] zkk]
9:
On chip, compute t t for i == 1,..,, 7i do
10:
11:
On chip, compute if mask k] A (1 < (i 1)) of then end for end for
12:
14:
Reduce l and all s; across threads on SRAM. for i = 1,
15:
do atomicAdd (2li, row', s.) atomicAdd + r,wl,, t t-s:) HBM.
16:
end for end for
18:
19:
end for
20: return
4 FlashMGLU: Efficient Kernel Implementation of MGLUs
Hardware-friendly functions such as attention has widespread application. Here we aim to make MGLUs efficient on modern hardware accelerators (GPUs) as well. In general, low batch size settings (for simplicity, we consider batch size 1) of LLMss generative inference require matrix-vector products. Such operations are almost always memory-bound, as the cost of reading the weight matrix dominates the compute time. A naïve Py PyToc implementation of MGLUs would issue nm separate matrix-vector multiplies (and corresponding element-wise Swish and multiplies), resulting in 7mm x 2 + full-precision weight memory reads per token-an an extremely inefficient pattern.
To address this, we implement fused CUDA kernel and a simple triton TTillet et 2019) kernel that:
Packs mask bits: Combine the nmm binary masks M, for each weight entry into a single 8-bit integer, so that one load fetches all mask information for block of weights.
Shared loading: For each thread block, load tile of the shared weight matrix W and the corresponding packed mask words in one coalesced transaction for all n.m x output vectors.
This design reduces the number of global memory reads from 7y x (W load) + (mask load) tile, while performing all Swish and element-wise multiplications in fast on-chip memory. As result, our kernel achieves up to 19.66x speedup over naïve PyTorch implementation on RTX5090 presents the detailed procedure, and Figure illustrates the forward
with nmm = 8. Algorithm |
5

---

pass. The core idea is to fetch all necessary weight and packed- packed-mask blocks for each output element in a single coalesced read and complete every arithmetic operation entirely in registers, thereby eliminating redundant global-memory traffic.
4.1 Discussion
Memory Load During Inference. Table 2 represents the parameter count and required memory load at inference time. Inference in modern transformer FFNs is typically bound by memory bandwidth rather than raw compute. In SwiGLU, the gating and value projections each require reading an FP16 weight matrix of size h X for a total of two matrices per token (memory load 16 x 2hd bits), whereas normal GELU actívation only requires one up-projection. SwiMGLU instead folds the two intermediate projections into a single FP16 matrix W E Raxd and applies 7.m binary masks M, E {0, h at inference to recover gating. The resulting per-token memory load is one FP16 matrix (16 x 2hd bits for up-projection) plus Nmm mask bits (nm hd bits). For T 1, the relative reduction
Table 2:
Parameter count and memory-access cost per token for the intermediate layers of FFN variants during FP16 inference. h and d denote the hidden and intermediate sizes, respectively.
16 2hd (16 hd + hd) 0.46875,
16 2hd
Le., up to 47% fewer bits transferred, directly translating to faster inference on memory- bound hardware. Mixtures with 7mm I 16 therefore still reduces the parameter count of a standard SwiGLU FFN at 16bit inference time, yet retain SwiGLU-Ievel expressivity during optimization
#Params #Params Memory Load (FP16)
Layer type
(Binary) (bits)
LU GLU MGLU
hd
0 Tom
16hd
P
2hd
32hd
hd
(16 + nmm. nm)) hd
Number of Parameters. SwiMGLU replaces
two full-rank SwiGLU weight matrices (2hd FP16 values) with one matrix plus ne masks, yielding a footprint of 16hd + nm hd bits at inference (masks stored in 1-bit form). During training, masks are kept as FP16 logits, adding nm hd FP16 parameters.
Computational Cost. At inference time (next token prediction), only the forward pass is exe- cuted. SwiMGLU incurs 2(1 + Tem) hd multiply-add operations per token, versus 6hd for SwiGLU. On memory- hardware the extra FLOPs are negligible: memory bandwidth dominates, and SwiMGLU's reduction in FP16 reads directly lowers inference latency. Training on the other hand is compute bound with the backward costing roughly twice the FLOPs of the forward. Masks are stored as FP16 logits and use straight-through estimation, adding nm hd FP16 parameters and one extra element-wing wise multiply-add per mask each pass. Consequently, the total training cost becomes (6 + 8nm) hd, compared to 18 hd per token in SwiGLU, introducing a runtime overhead according to the mask size.
5 Experiments
We demonstrate the effectiveness and efficiency of SwiMGLU. We selected the Llama 3 architecture (Grattafiori et al., 2024) with SwiGLU as the baseline and compare the different FFN layers.
5.1
Setup
Model Configuration. Our baseline model follows the Llama family design (Grattafiori et al. 2024). We train models at two scales: a 159M small (12 layers, h = 768, d = 3072) and a 1.08B large (16 layers, h = 2048, d = 8192. For SwiMGLU, we substitute each SwiGLU layer with the SwiMGLU layer, keeping all other architectural components (e.g. attention, normalization, embeddings) identical. Note that SwiMGLU would have a smaller model size as the number of projection weights are reduced.
Hyperparameters. All model are trained using the AdamW optimizer (3: = 0.9, 32 0.99, c= 1 x 10-3) 8) with a learning rate of 3 x 10-4 and 1 x 10-4 for small and large models respectively. weight decay of 1, linear warmup for the first 10% steps, followed by cosine decay. full list of hyperparameters and training settings are listed in Appendix
For Llama- IB (h- = 2048, d= 8192) the FFN weights per layer shrink from 96MB to 64MB, while the binary mask costs only 2MB in boolean form.

---

Table Zero-shot accuracy (%) on downstream tasks and validation perplexity. #weights represent
the number of weight parameters excluding masks bits.
7lm
#weights
PPL.
AreEt
AreCt
HS*
PiQAt
SciQt
WGt
Avg
GELU
113M
25.8
47.47
19,45
28.09
60.61
64.80
52.41
45.47
SwiGLU
141M
23.7
48.15
20.05
28.53
61.43
67.90
51.14
46.20
SwiMGLU
113M
25.0
48.91
19.28
28.25
60.72
69.00
50.20
46,06
SwiMGLU
2
113M
24.5
49.12
19.97
28.49
60.01
70.60
51.53
46.62
SwiMGLU
113M
23.9
48.99
20.56
28.49
61.70
69.10
50.04
46.48
SwiMGLU
113M
23.5
48.65
20.56
28.63
61,53
68.00
51.54
46.49
SwiGLU
1.08B
12.3
64.94
28,92
37.20
69.15
84.50
51.30
56.00
SwiMGLU
808M
13.0
63.72
27.73
36.20
68.61
83.00
54.30
55.59
SwiMGLU
2
808M
12.7
62.08
26.71
36.52
68.44
84,20
51.85
54.97
SwiMGLU
808M
12.4
65.78
28.92
37.69
69.26
84.20
55.25
56.85
Table 4: Two-shot accuracy (%) on downstream tasks and validation perplexity. #weights represent the number of weight parameters excluding masks bits.
#weights
PPL
ArcEt
ArcCt
HS*
PIQAt
SciQt
WG1
Avg
GELU
113M
25.8
48.57
19.45
27.87
60.77
61.10
52.88
45.11
SwiGLU
141M
23.7
48.65
20.14
28.64
61.70
62.30
51.70
45.52
SwiMGLU
113M
25.0
48.48
19.03
28.16
60.45
61.40
50.75
44.71
SwiMGLU
113M
24.5
49.62
21.42
28.55
60.39
62.50
51.14
45.60
SwiMGLU
113M
23.9
49.41
21.16
28.61
62.51
66.10
50.59
46.40
SwiMGLU
8
113M
23.5
52.02
19.62
28.54
62.08
66,80
51.38
46.74
SwiGLU
1.08B
12.3
66.16
31.48
37.25
69.21
88.30
51.78
57.36
SwiMGLU
808M
13.0
65.57
30.20
36.25
69.59
88.50
53.83
57.32
SwiMGLU 808M
12.7
65.07
31.66
36.32
68.77
88.70
52.17
57.11
SwiMGLU
808M
12.4
67.13
30.55
37.53
69.48
88.60
53.91
57.87
Datasets and Metries. We pre-train both baseline and SwiMGLU models the Feb-Web-Edu 100B dataset (Penedo et al. 2024) with small models being trained on a 10B token subset. For downstream evaluation, we report zero-shot and two-shot accuracy on six standard benchmarks: ARC Easy (ArcE) (Clark et al. 2018). ARC Challenge (AreC) (Clark et al.) 2018), HellaSwag (HS) (Zellers Bisk Bisk al. (2020). SciQ (Welbl al. (017)). and Winogrande (WG) (Sakaguchi
al. 2019), PiQA
ct al. 2021). We utilize the LM Evaluation Harness (Gao et al. 2024) for standardized performance evaluation.
5.2 Downstream Performance
Downstream Accuracy. Tables 3 and present the zero-shot and two-shot performance respectively. We observe that SwiMGLU generally matches or outperforms SwiGLU when 4 or more masks are employed. For example
showcases average of 56.85% compared to SwiGLU's 56.00%
These results highlight the superior performance of model using the SwiMGLU activation function. Increasing the mask count generally improves the model performance.
Latency. Figure 4 shows the wall-clock la- tency of single projection layer in a single batch setting measured on an RTX 5090 at FP16 precision. We evaluate three imple- mentations: our fused CUDA kernel (Flash- MGLU), a Triton re-implementation, and the naïve PyTorch GLU nn Linear baseline.
Latency (ms)
Latency (ms)
[ajh 2048 d 8192
0b/A 4006 d 14338
For the large setting (h= 2048, d= 8192) with R.n 8, FlashMGLU cuts the ma- trix-vector time from 0.5210 ms to 0.0834
0.01
0.01
ms, delivering a 19.66x speed-up over the PyTorch baseline and 6.25x acceleration even on our triton prototype. Importantly, these gains persist at much larger scales: with
Number
Number
(PyTa))))
MGLU (PyTorch)
Flash
Flachii (CUDA)
Figure Latency comparison

---

Training Loss
Training Perplexity
Training Loss
Training Perplexity
30
3.8
3.6
34
3.4
(M
0
10
10
10
10
Processed Tokens
Processed Tokens (8)
Processed Tokens (9)
Processed Tokens 18)
GELU
SWIGLU
SwiMGLU (5)
16
2,0
3.0
2.8
2.8
14
2.6
2.6
0
20
40
60
30
20
40
60
80
20
40
60
50
20 40
60 50
Processed Tokens
Processed Tokens (8)
Processed Tokens (8)
Processed Tokens 18)
SwIGGLU
Sw MGLU (Po m4
-2
Figure 5: Comparison of learning curves for different FFN architectures. The top and bottom rows illustrate the changes in training loss / training perplexity of small and large models respectively. The left columns compare existing methods against SwiMGLU, and the right columns compare the number of masks.
ArcE Accuracy (%)
HS Accuracy (%)
PIQA Accuracy (%)
Average Accuracy (%)
28
35
40
50
4.1
60
60
40
60
Processed
SWIGLU
SWiMGLU -1)
SwiMGLU [Rgg
SwiMGLU [e
Figure 6: Comparison of downstream task scores across different FFN architectures large models. In all metrics, the proposed method, SwiMGLU nm = 4 achieves the best performance.
larger intermediate sizes (h 4096, 14336) used in larger Llama-3.2 8B models, our optimized CUDA kernel still outpaces the PyTorch baseline by an impressive 18.0x, and a 11.1x speed-up in our triton prototype.
While the standard GLU implementation must fetch two full-precision weight matrices from global memory, FlashMGLU accesses each weight exactly once and performs all mask operations directly within on-chip registers. By fusing weight loading with mask evaluation, FlashMGLU significantly increases arithmetic intensity, effectively saturating the GPU's Streaming Multiprocessor pipeline and reducing latency by up to 1.51x under the single-mask variant. As a result, our native CUDA implementation of FlashMGLU consistently outperforms both the naïve MGLU and the standard GLU baselines. Although Triton incurs slightly higher overhead on very small matrices, this gap narrows as the hidden dimensions increase. Moreover, since CUDA allows developers to explicitly control memory access patterns, native CUDA code achieves higher performance overall. A full head-to-hhe to- head comparison between FlashMGLU and highly tuned standard GLU kernel presented in Appendix E
Scaling with the number of masks. Under naive PyTorch MGLU kernel, runtime scales nearly linearly for 7m < 8, doubling the number of mask mixtures almost doubles the latency, driven by proportional increases in memory-trafie loading strategy decouples execution time from mask count, yielding only marginal slowdowns as nm increases. Even under computationally intense eight mask variant, we still observe 1.15x larger variants. Once = 8, register-file
In contrast, FlashMGLU's unified weight-and-ng mask
speed-up on small models and 1.24x acceleration pressure begins to spill into local memory, and the performance advantage over the standard GLU kernel gradually crodes, which is effect that is especially noticeable in Triton prototype.

---

5.3 Ablations and Analysis
Training Dynamics. Figure([plott token-level cross- pross-entropy (left axis) and perplexity (right axis) over the first 150k optimization steps for both small and large models trained with GELU, SwiGLU, and our SwiMGLU variants. Models using SwiMGLU with 7mm = 4 consistently show comparable or lower losses compared to those using SwiGLU. fillstttesss the downstream performance on ARC-Easy, HellaSwag, PiQA, and the average accuracy across all downstream tasks. SwiMGLU with Rmm outperforms SwiGLU on all tasks, with notable improvements, demonstrating strong generalization capabilities of MoGE despite its reduced parameter count.
Training Cost. While SwiMGLU significantly reduces seereence-time memory access and latency, incurs a higher computational cost during pretraining. As shown in Appendix A (Table ( achieving comparable downstream performance SwiGLU requires upproximately twice the wall- clock training time due to the additional optimization of mask parameters. This trade-off reflects a prioritize lower inference cost and memory footprint at the expense of moderately higher training time. We believe this is a reasonable compromise, as real-world deployments of LLMS are typically dominated by inference workloads rather than pretraining.
deliherate design choice:
Training Stability. At default learning rates, both MGLU and SwiGLU converge with comparable smoothness. However, under higher learning rates, where SwiGLU experiences pronounced loss spikes, the masked variants effectively dampen these excursions and recover more rapidly. The four-mask configuration, in particular, proves exceptionally robust. Detailed loss curves for all learning-rat -aat settings are provided in Appendix B
Number of Masks mm For the small model size, increasing the number of masks from 7m = 1 to 7m = 8 results in strictly mono- tonic reduction in training loss. Concurrently, this also improves the two-shot average ac- curacy, which rises from 44.71% to 46. 46.74% Nevertheless, the marginal benefit diminishes once R'm = 8, while memory and compute overhead continue to increase. Empirically, the n.m = setting achieves the optimal balance between quality and resource usage. These findings suggest that a compact col- lection of complementary masks can fully recover, and occasionally surpass, the expressive power of a dense SwiGLU gating mechanism.
Training
Training Perplexity
30
3.4
3.2
8
10
Processed Tokens
Processed Tokens (B)
Learned
Fixed
Figure 7: Training loss and perplexity of learned vs.
fixed masks in small configuration and n.m 2.
Learned vs. Fixed Masks. Figure 7 contrasts models in which the masks are co-optimized with the network parameters against models that rely on randomly sampled but fixed masks. Allowing the masks learn consistently tracks lower training loss throughout and improves the final perplexity by roughly 0.5. By comparison, freezing the masks essentially collapses the learning curve to that of the single-mask (nom 1) bascline, underlining that it is alone-that drives the gain in expressivity. Complete results are provided in Appendix
learning not sparsity
Mask Distribution, Inspection of the learned binary masks shows that their activation ratio is not foreibly balanced: depending on the layer, the proportion of ones can skew either above or below 50% (typically ranging from roughly 45% to 55%). Full statistics for every layer and model size are provided in Appendix C This flexibility lets the network learn how much capacity to devote to the gate versus the value pathway, effectively tuning the gate-value trade-off on per-layer basis. Because a dense SwiGLU block lacks such adaptive capacity allocation, this mechanism offers a plausible explanation for the cases in which MGLU surpasses SwiGLU in accuracy.
6 Conclusion and Future Work
In this work, we introduced Masked Gated Linear Units (MGLUs), a novel family of ford-forward activations that recover the expressivity of traditional GLUs using a single shared weight matrix sculpted by learnable binary masks. Our Mixture of Element-wise Gating (MoEG) design not only matches or exceeds the performance of the commonly used SwiGLU activation on a variety of finguage-understangg benchmarks, but also delivers substantial efficiency gains at inference time.

---

In terms of limitations, MGLUs are more computationally expensive than other GLU variants, since they increase arithmetic complexity exchange for a reduced memory footprint. However, by developing FlashMGLJU, an optimized kernel that fuses mask unpacking with the core matrix-vector operations, we achieve up to a 19.66x speedup over naïve PyTorch MGLU implementation on modern GPUs, while reducing memory-bandwidth requirements and latency compared to standard GLU variants. These advances pave the way for richer gating mechanisms to be deployed in latency- and memory-constrained settings without compromising model quality.
7 Acknowledgement
This work was supported by DENSO IT LAB Recognition, Control, and Learning Algorithm Collaborative Research Chair (Science Tokyo).
This work used computational resources TSUBAME4.0 supercomputer provided by Institute of Science Tokyo through the HPCI System Research Project (Project ID: hp240170).
References
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proc. Annual Conference Neural Information Processing Systems (NeurIPS). 2017.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, and Amanda et al. Askell. Language models are few-shot learners. In Proc. Annual Conference on Neural Information Processing Systems (NeurIPS), 2020.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, brian ichter, Fei Xia, Ed H. Chi, Quoc V Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. In Proc. Annual Conference on Neural Information Processing Systems (NeurIPS), 2022.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, and Alex Gray al. Training language models to follow instructions with human feedback. In Proc. Annual Conference on Neural Information Processing Systems (NeurIPS), 2022.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.2178. 2024.
OpenAI. Gpt-4 technical report. arXiv preprint arXiv:2303.0877 2024.
Gemma Team. Gemma technical report. arXiv preprint arXivaarXiv:2503.11975 2025.
Amir Gholami, Zhewei Yao, Sehoon Kim, Coleman Hooper, Michael W. Mahoney, and Kurt Keutzer. Ai and memory wall. arXiv preprint arXiv2403.1411 2024.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing by generative pre-training. OpenAl Technical Report, 2018
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher R'e. FlashAttention: Fast and effoee-e exact attention with io-awareness. In Proc. Annual Conference on Neural Information Processing Systems (NeurIPS), 2022.
Tri Dao. FlashAttentionn-2: Faster attention with better parallelism and work partitioning. In Proc. International Conference on Learning Representations (ICLR), 2024.
Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. FlashAttention-3: Fast and accurate attention with asynchrony and low-precision. In Proc. Annual Conference on Neural Information Processing Systems (NeurIPS), 2024.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In Proc. Conference Language Modeling (COLM), 2025.
10

---

Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. In Proc. International Conference on Machine Learning (ICML), 2024.
Vinod Nair and Geoffrey E. Hinton. Rectified linear units improve restricted boltzmann machines. In Proc. International Conference on Machine Learning (ICML), pages 807-814, 2010.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint
arXiv: 08415, 2016.
Yano N. Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In Proc. International Conference Machine Learning (ICML), pages 933-941, 2017a.
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202, 2020.
Prajit Ramachandran, Barret Zoph, and Quoe V. Le. Searching for activation functions. arXiv preprint arXiv:1710.05941, 2017.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In Proc. International Conference on Learning Representations (ICLR), 2019.
Mengzhou Xia, Zexuan Zhong, and Danqi Chen. Structured pruning learns compact and accurate models. In Proc. Annual Meeting of the Association for Computational Linguistics (ACL), 2022.
Bairu Hou, Qibin Chen, Jianyu Wang, Guoli Yin, Chong Wang, Nan Du, Ruoming Pang, Shiyu Chang, and Tao Lei. costuuttion-oolloing pruning for large language models. In Proc. International Conference on Machine Learning (ICML), 2025.
Fabrizio Sandri, Elia Cunegatti, and Giovanni Iacca. 2SSP: A two-stage framework for structured pruning of lims. arXiv preprint arXiv2501.17771 2025.
Saleh Ashkboos, Maximilian L. Croci, Marcelo Gennari do Nascimento, Torsten Hoefler, and James Hensman. SliceGPT: Compress large language models by deleting rows and columns. In Proc. International Conference on Learning Representations (ICLR), 2024.
Gongfan Fang, Hongxu Yin, Saurav Muralidharan, Greg Heinrich, Jeff Pool, Jan Kautz, Pavlo Molchanov, and Xinchao Wang. MaskLLM: Learnable semi-structured sparsity for large language models. In Proc. Annual Conference on Neural Information Processing Systems (NeurIPS), 2024.
Mingjie Sun, Zhuang Liu, Anna Bair, and J. Zico Kolter. A simple and effective pruning approach for large language models. In Proc. International Conference on Learning Representations (ICLR), 2024.
Elias Frantar and Dan Alistarh. SparseGPT: massive language models can be accurately pruned in one-shot. In Proc. International Conference on Machine Learning (ICML), 2023.
Zhenyu Zhang, Zechun Liu, Yuandong Tian, Harshit Khaitan, Zhangyang Wang, and Steven Li. R-sparse: Rank-aware activation sparsity for efficient LLM inference. In Proc. International Conference on Learning Representations (ICLR), 2025.
Daniel HAZIZA, Timothy Chou, Dhruv Choudhary, Jesse Cai, Luca Wehrstedt, Francisco Massa, Jiecao Yu, Geonhwa Jeong, Supriya Rao, and Patrick Labatut. Accelerating transformer inference and training with 2:4 activation sparsity. In Sparsity in LLMs (SLLM); Deep Dive into Mixture of Experts, Quantization, Hardware, and Inference, 2025.
Seyed Iman Mirzadeh, Keivan Alizadeh-Vahie Sachin Mehta, Carlo C del Mundo, Oncel Tuzel, Golnoosh Samei, Mohammad Rastegari, and Mehrdad Farajtabar. ReLU strikes back: Exploiting activation sparsity in large language models. In Proc. International Conference Learning Representations (ICLR), 2024.
Zichang Liu, Jue Wang. Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivastava, Ce Zhang, Yuandong Tian, Christopher Re, and Beidi Chen. Deja vu: Contextual sparsity for efficient LLMS inference time. In Proc. International Conference on Machine Learning (ICMI.) 2023.
11

---

James Liu, Pragaash Ponnusamy, Tianle Cai, Han Guo, Yoon Kim, and Ben Athiwaratkun. Training- free activation sparsity in large language models. In Proc. International Conference on Learning Representations (ICIRR) 2025.
Yann N. Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In Proc. International Conference on Machine Learning (ICML), 2017b. 7b.
Yoshua Bengio, Nicholas L'eonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint
2013.
Robert Jacobs, Michael Jordan, Steven J. Nowlan, and Geoffrey Hinton, Adaptive mixtures of local experts. Neural Comput., 3(1):79-87, 1991.
Michael I. Jordan and Robert A. Jacobs. Hierarchical mixtures of experts and the EM algorithm. Neural Comput., 6(2):181 11-214,, 1994.
Philippe Tillet, H. T. Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proc. ACM SIGPLAN International Workshop Machine Learning and Programming Languages (MAPL), page 10-19, 2019.
Guilherme Penedo, Hynek Kydlíček, Loubna Ben allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, and Thomas Wolf. The fineweb datasets: Decanting the web for the finest text data at scale. In Proc. Annual Conference on Neural Information Processing Systems (NeurIPS), 2024.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv: 1803:0444 2018.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. HellaSwag: Can a machine really finish your sentence? In Proc. Annual Meeting of the Association for Computational Linguistics (ACL), 2019.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language. In Proc. AAAI Conference Artificial Intelligence, 2020.
Johannes Welbl, Nelson F. Liu, and Matt Gardner. Crowdsourcing multiple choice science questions. In Proc. Workshop on Noisy User-generaedd Text, 2017.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi Winogrande: adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99-106, 2021.
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding. Jeffrey Hsu, Alain Le Noac h. Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. The language model evaluation harness. 2024. doi: 10.528 1022///ooooo.. 12608602.
Kazuki Fujii, Taishi Nakamura, and Rio Yokota. llm-recipes, May 2024. URL https://t//thbbb com/ -rricip
William Fedus, Barret Zoph, and Noam M. Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. J. Mach. Learn. Res., 23:120:1 -120:39,, 2021.
12

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justifiicton right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No] ", it is perfectly acceptable to answer "[No] provided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: This submission introduces novel MGLU FFN architecture in Large Lan- guage Models, including fast FlashMGLU kernel for inference.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made 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 is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
13

---

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

---

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 example, if the contribution is a 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 it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open- source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibilityy 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
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification:
The code will be released with permissive license in the near future.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
wwt.:/owbbbiiiooooooooooooo for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a 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 CodeeSumiiii wwwiiisssooosiiiiio
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 are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
15

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We discuss datasets and syper-parametent used for training and testing in the appendix.
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: Error bars are not available in this submission.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time exccution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We report the number of GPUs and GPU hours in appendix.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
16

---

The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics htttps:/neurrps.. http://wwwpp.ccoopuu..i.caaaceae
Answer: [Yes]
Justification: This submission follows the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. if there special consid- eration 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: [Yes]
Justification:
This submission discusses potential societal impacts in the appendix.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a 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 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 models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification: This paper does not include publicly released models nor datasets.
Guidelines:
The answer NA means that the paper poses such risks.
17

---

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 to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards 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
Are the creators 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?
Question:
Answer: [Yes]
Justification: All assets are properly credited.
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 of 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 a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not include new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent 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 crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
18

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed 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: The paper does not involve crowdsourcing nor research with 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 in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods this research? Note that if the L.M used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration not required.
Answer: [NA]
Justification: LLMs do not involve in method development in the paper.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our L.M policy ottt:://eeerrrp.cc//ooffrrrrrr/ for what should or should not be described.
19

---

A
Experimental Details
Here we provide more details about the model architecture, training configurations and resources used in our experiments.
A.1 Model Architecture and Training Configuration
Table summarizes the architectural configurations shared across all experiments. Table lists the number of weight and mask parameters, along with the corresponding estimated storage sizes at inference time for each model variant.
All experiments reported in this paper are implemented based on the 11m-recipes framework (Fuji al. 2024). All models are trained on TSUBAME supercomputer with NVIDIA H100 GPUs (94GB), with small models trained on 4GPUs and large models trained on 16GPUs with Fully Sharded Data Parallel (FSDP). Training GPU hours shown in Table
Table 5: Model architecture of small and large variants.
Model Size
h
d
Context Length
#Heads
#Layers
small
768
3072
1024
24
12
large
2048
8192
4096
32
16
Table 6: Number of weight and mask parameters, estimated storage size, and training GPU hours for each model configuration.
Scale
Model
7mm
#Weights
#Masks
Size (MB)
GPU Hours
GELU
113M
215
18
SwiGLU
141M
0
269
22
SwiMGLU
113M
28.3M
219
20
small
SwiMGLU
113M
56.6M
222
24
SwiMGLU
4
113M
113M
229
33
SwiMGLU
8
113M
226M
242
52
SwiGLU
L.08B
2052
768
SwiMGLU
808M
268M
1573
752
large
SwiMGLU
808M
537M
1605
1008
SwiMGLU
4
808M
1.07B
1669
1376
A.2 Hyperparameters
Table 7 lists the hyperparameters that we use by default at training time for all our experiments.
Table 7: Pretraining hyperparameters for small and large models.
small
large
Optimizer
AdamW
AdamW
Learning Rate (LR)
3E-4
1E-4
Minimum LR
3E-5
IE-5
LRR Schedule
cosine
cosine
Weight Decay
0.1
0.1
81
0.9
0.9
82
0.99
0.99
1E-8
IE-8
Gradient Clipping
Global Batch Size
512
512
Warmup Steps
1000
4800
20

---

Training Loss
Training Perplexity
30
3.6
3.4
25
3.2
2
10
0
2
6
8
10
Processed Tokens (B)
Processed Tokens (B)
ReMGLU
GeMGLU
SwiMGLU
Figure 8:
Trainig Curves of small MGLU variants with different activation functions. Left:
training loss; right: validation perplexity.
Table 8: shoo accuracy (%) on downstream tasks and validation perplexity across different
activation functions. #weights represent the number of weight parameters excluding masks bits.
7m
#weights
PPL.
ArcEt
AreC1
HS*
PiQAt
SciQt
WG1
Avg
SwiGLU
141M
23.7
48.15
20.05
28.53
61.43
67.90
51.14
ReMGLU
113M
24.3
49.66
19.71
28.45
61.53
69.10
49.80
46.38
GeMGLU
4
113M
24.0
48.65
19.54
28.50
62.30
69.30
52.17
46.74
SwiMGLU
113M
23.9
48.99
20.56
28.49
61.70
69.10
50.04
46.48
Table 9: Two-shot accuracy (%) downstream tasks and validation perplexity across different activation functions. #weights represent the number weight parameters excluding masks bits.
7m
#weights
PPL.
AreEt
AreCt
HSC
PiQA
SciQt
WGt
Avg
SwiGLU
141M
23.7
48.65
20.14
28.64
61.70
62.30
51.70
45.52
ReMGLU
113M
24.3
49.28
19.71
28.25
61.32
61.80
51.22
45.26
GeMGLU
113M
24.0
49.28
20.56
28.74
62.13
66.40
52.25
46,56
SwiMGLU
113M
23.9
49.41
21.16
28.61
62.51
66.10
50.59
46.40
B Additional Experiments
B.1 Other Activation Functions.
In the main paper we concentrated on SwiMGLU, whose gating function relies on the Swish non- linearity used in standard SwiGLU. To verify that the advantages of our mask-based design not tied to a single activation, this appendix evaluates several alternatives that are widely adopted in large-scale language models- GELU, ReLU, and SiLU. For each activation replace the gating function g(-) in Eq. (3) while keeping all other architectural choices and training identical to the baseline. We report both pre- pre-training perplexity and downstream accuracy of small models so that we can directly compare the impact of each non-linearity under the same experimental conditions.
Training Curve. Figure&plots training loss and perplexity for the various activations. ReMGLU converges more slowly than both GeMGLU and SwiMGLU, with SwiMGLU achieving the fastest convergence and slightly outperforming GeMGLU.
Downstream Evaluation. Table & and Table show the downstream evaluation scores on different activation functions in MGLU. The masked design improves or matches the SwiGLU baseline across all activations while using fewer trained weights, confirming that the capacity unlocked by mask pairs is activation-agnostic. Smoother nonlinearities (GELU and SiLU) benefit most: GeMGLU achieves the best average zero-shot accuracy (46.74%) and ties for the best two-shot average (46.56%) whereas SwiMGLU attains the lowest validation perplexity (23.9). ReMGLU, though slightly weaker than its smoother counterparts, and comparable the SwiGLU baseline average accuracy.
21

---

1%
x
++
Swish
M&W
Mow
M&w
McW
MgoW
M.W
MeW
Router
Figure 9: Diagram of Top-1 SwiMGLU block. We illustrate one token being routed across four mask experts, where the router independently routes each token. The Top-1 SwiMGLU layer returns the output of the selected experts multiplied by the router gate value.
B.2 Top-K Routing
Summation over every masking route maximizes capacity but may be wasteful for tokens that require only subset of specialized subspaces. Inspired by routing techniques in switch transformers (Fedus) et al. 2021), we therefore explore Top-K routing variant of MGLU that, at inference time, activates only the K masks whose logits yield the highest gate magnitudes. The subsection first details the routing algorithm and its lightweight implementation iicemennttion--rquiinngg single additional linear router over the input features, then evaluates downstream scores K varies from 1 (fully sparse) to n'm (fully dense). We show that, on small configuration, routing with K = 2 retains most of the accuracy gains of full MGLU with n. 4, while reducing computational FLOPs by 2x, offering practical knob for deployments with stringent latency budgets.
Top-K routed MGLU. Given input token representation R lightweight "router" computes a vector of logits
i xW,
1m
where W, is a learned h X Tem matrix. Following Fedus et al. (2021), retain only the K largest logits. Applying a softmax over this truncated vector yields sparse gating weights
G(x) Softmax (TopK(e))
Rm
(5)
The Top-K routed MGLU is denoted by
Am
MGLUTg..<())
Gl)) Σ
@ a(Mow).
(6)
Because G(x) contains most K non-zero terms, only those K masked projections need be evaluated-rrducing memory traffic and latency compared with summing over all n.m routes while preserving most of the accuracy gains of full MGLU. Figure illustrates the overall architecture.
Training Curves. Figure I0 presents the training loss and perplexity across different routing configu- rations. We observe that both Top-4 and Top-2 routed SwiMGLU closely track the loss trajectory of the non-routed SwiMGLU, indicating that sparse routing with multiple active masks retains most of the model's learning capacity. In contrast, Top-1 routing introduces noticeable performance degradation, suggesting that activating only a single expert per token limits representational power.
Downstream Evaluation. Table IO]and Table II ssmmmarize validation perplexity and task perfor- mance for each routing strategy. The Top-4 router (K = 4) consistently dominates, delivering the lowest perplexity in both settings, and the highest average accuracy. The Top-2 variant (K 2) offers a favorable trade-off: matches or surpasses the unrouted SwiMGLU (nn 44) on most metrics, boosts average zero accuracy 46.64 %, and retains competitive two-shot score of 45.79 %.
22

---

Training Loss
Training Loss
3.6
3.6
3.4
3.4
3.2
3.2
10
10
Processed Tokens (B)
Processed Tokens (B)
Training Perplexity
Training Perplexity
30
30
25
25
6
8
10
0
2
4
8
10
Processed Tokens (B)
Processed Tokens (B)
Top
Top 2
Top
Top 2
Sum
Figure 10: Learning curves of small SwiMGLU models under different Top- routing strategies. Left: training loss and validation perplexity for K E {1,2, 4}; right: Top-2 routing compared with the non-routed SwiMGLU baseline.
Table 10: Zero-shot accuracy (%) downstream tasks and validation perplexity across different routing coefficient K. #weights represent the number of weight and router parameters excluding masks bits. The boldface and underline indicate, respectively, the best and second-best value per
column.
K
#weights
PPL
ArcEt
AreCt
HSt
PiQAt
SciQr
WG*
Avg t
SwiGLU
141M
23.7
48.15
20.05
28.53
61.43
67.90
51.14
46.20
SwiMGLU
113M
23.9
48.99
20.56
28.49
61.70
69.10
50.04
46.48
SWIMGLU
1
113M
25.2
48.11
19.80
27.89
61.04
68.90
50.67
46.07
SwiMGLU
4
2
113M
24.0
48.86
20.56
28.29
60.55
69.50
52,09
46.64
SwiMGLU
4
4
113M
23.8
49.20
21.84
28.70
61.37
71.80
53.43
47.72
Table 11: Two-shot accuracy (%) on downstream tasks and validation perplexity across different
routing coefficient K. #weights represent the number of weight and router parameters excluding mask bits. The boldface and underline indicate, respectively. the best and second-bes value per column.
Tmm K
#weights
PPLÍ
ArcEt
AreC1
HS²
PiQAt
SciQr
WG1
Avg t
SwiGLU
141M
23.7
48.65
20.14
28.64
61.70
62.30
51.70
45.52
SwiMGLU
4
113M
23.9
49.41
21.16
28.61
62.51
66.10
50.59
46.40
SwiMGLU
4
1
113M
25.2
49.03
20.31
27.95
60.50
59.90
52.17
44.98
SwiMGLU
4
2
113M
24.0
51.39
20.65
28.71
62.08
60.40
51.54
45.79
SwiMGLU
113M
23.8
49.79
22.27
28.82
62.13
63.00
52.57
46.43
contrast, Top-1 routing (K 1) shows significant drop in both perplexity and average accuracy, suggesting that allocating at least two experts per token is critical for maintaining representation power. Overall, these results indicate that 2 not only reduces compute but also enhances generalization across a diverse suite of downstream tasks.
23

---

Training Loss
Training Loss
=1
3.6
3.6
3.4
3.4
3.2
3.2
10
10
Processed Tokens (B)
Processed Tokens (B)
Training Perplexity
Training Perplexity
30
30
12m
Fim
25
25
0
8
10
0
2
4
8
10
Processed Tokens (B)
Processed Tokens (B)
Learned
Fixed
Figure 11: Training curves of learned vs. fixed masks in small SwIMGLU models. Lef: = 1; Right: nn = 2.
Table 12: Zero-shot accuracy (%) downstream tasks and validation perplexity small models.
Tbm
Mask Type
PPL.
ArcE1
ArcC1
HS1
PiQA
SciQt
WG1
Avg f
SwiMGLU
Learned
25.0
48.91
19.28
28.25
60.72
69.00
50.20
46.06
SwiMGLU 1
Fixed
25.1
47.60
21.08
28.19
61.10
68.30
51.54
46.30
SwiMGLU
Learned
24.5
49.12
19.97
28.49
60.01
70.60
51.53
46.62
SwiMGLU
Fixed
25.1
48.32
28.23
61.81
67.30
50.36
45.84
Table 13: Two-shot accuracy (%) on downstream tasks and validation perplexity on small models.
Th
Mask Type
PPI.
ArcEt
ArcCt
HS*
PiQAt
SciQt
WG
Avg f
SwiMGLU
Learned
25.0
48.48
19.03
28.16
60.45
61.40
50.75
44.71
SwiMGLU
Fixed
25.1
48.11
20.14
28.02
61.21
60.30
51.78
44.92
SwiMGLU
2
Learned
24.5
49.62
21.42
28.55
60.39
62.50
51.14
45.60
SwiMGLU
2
Fixed
25.1
47.14
20.31
28.24
61.43
61.40
51.93
45.07
B.3 Learned vs. Fixed Masks.
This section compares learned and fixed masks in small SwiMGLU models. While previous sections focus on the number and structure masks, here we investigate the benefit of co-optimizing masks alongside model weights.
Training Curve. Figure II shows the training loss and perplexity when using learned versus fixed masks. While with n,, = 2, learned masks consistently outperform fixed ones, with n.m the
difference is subtle.
Downstream Evaluation. Table Table and Table shows the downstream task scores across different
mask configurations. The learned mask configuration with n 2 also improves downstream accuracy compared to fixed masks, highlighting the importance of mask combination adaptation. These results suggest that expressivity not solely due to sparsity, but also driven by mask learning with multiple masks.

---

Training Loss
Training Perplexity
30
3.6
3.4
25
3.2
10
10
Processed Tokens (B)
Processed Tokens (B)
Tho=
Rum =2
Rin =4
Tim
Tkm = 16
Figure 12: Training curves of small SwiMGLU models across different mask count.
Table 14: Zero-shot accuracy (%) on downstream tasks and validation perplexity. #weights represent the number of weight parameters excluding masks bits. The boldface and underline indicate, respec- tively, the best and second-best value per column.
7m
#weights
PPL4
ArcE1
AreC1
HS
PiQAt
SciQt
WG1
Avg
SwiMGLU
113M
25.0
48.91
19.28
28.25
60.72
69.00
50.20
46.
SwiMGLU
2
113M
24.5
49.12
19.97
28.49
60.01
70.60
51.53
46.62
SwiMGLU
4
113M
23.9
48.99
20.56
28.49
61.70
69.10
50.04
46.48
SwiMGLU
113M
23.5
48.65
20.56
28.63
61.53
68.00
51.54
46.49
SwiMGLU
113M
23.3
48.15
21.08
28.63
61.59
68.50
50.91
46.47
Table 15: Two-shot accuracy (%) on downstream tasks and validation perplexity. #weights represent the number of weight parameters excluding masks bits. The boldface and underline indicate, respec- tively, the best and second-best value per column.
7m
#weights
PPL4
ArcEt
AreCt
HS
PiQAt
SciQt
WG1
Avg
SwiMGLU
113M
25.0
48.48
19.03
28.16
60.45
61.40
50.75
44.71
SwiMGLU
2
113M
24.5
49.62
21.42
28.55
60.39
62.50
51.14
45.60
SwiMGLU
113M
23.9
49.41
21.16
28.61
62.51
66.10
50.59
46.40
SwiMGLU
113M
23.5
52.02
19.62
28.54
62.08
51.38
46.74
SwiMGLU
113M
23.3
50.76
21.08
28.66
61.70
66.90
$2.25
46,89
B.4
Scaling n.m 16
We investigate the impact of increasing the number masks 1.m up to 16. As shown in Figure 12 scaling 7hm leads steady improvements in both training loss and perplexity. However, the benefit diminishes beyond n = 4, and the gains from nmm = 8 to n.m = 16 are marginal.
Downstream Evaluation. Tables 14 and 15 report zero-shot and two-shot accuracy across different values of ny. The n.
setting achieves the best average accuracy, but the improvements over 7.m = 4 and n.m. are small. These results suggest that increasing the number of complementary masks helps, but overly large mask sets yield diminishing returns.
B.5 Partial Mask Ablation
In the standard single-mask MGLU layer (Eq. 2), the gate and value streams are computed from complementary subspaces of shared weight matrix W defined by binary masks M and M. To assess the necessity of these masks, introduce three ablation variants that remove the masks from
25

---

Training Loss
Training Perplexity
4.4
4.0
3.6
30
3.2
26
0
2
4
3
Processed Tokens (B)
Processed Tokens (B)
SwiMGLU (Ray
No Gate Mask
No Value Mask
No Masks
Figure 13: Training loss and perplexity for mask-ablation variants. All ablation variants- Gate Mask, No Value Mask, and No Masks Wasks-converree to higher loss and perplexity compared to the fully masked MGLU baseline, highlighting the necessity of maintaining complementary mask- defined subspaces.
gate, value, or both streams:
No Gate Mask (Dense Gate): hNG(œ) = (xW) @ (a(M o W)),
(7)
No Value Mask (Dense Value):
hxv(2) = g(x(M @ w)) @ (aW),
(8)
No Masks (Fully Shared): hxM(() = g(*W) @ (xW).
(9)
These variants maintain the multiplicative gate-value interaction characteristic of the MGLU architec- ture while systematically testing the significance of the mask-defined complementary subspaces.
Training Curves. Figure 13 compares training loss and validation perplexity these mask-ablation variants in the small model setting. All ablation variants (No Gate Mask, No Value Mask, and No Masks) show slower convergence and higher final loss compared to the fully masked baseline. These results highlight the importance of distinct, complementary subspaces defined by masks for optimal MGLU performance.
C
Mask Distribution
To understand how the model allocates capacity when multiple masks are learned, we measure the fraction of rows in the shared projection matrix W that each mask devotes to the gate pathway (higher values indicate more gate parameters; the remainder are routed the value pathway). Figure 14 plots this layer-wise gate ratio for models trained with n.m. 1,2, and 4.
All configurations exhibit shallow U-shape: gate capacity is largest in the first layer, reaches minimum around the middle of the network, and rises again toward the top. Because all masks share the same underlying weight matrix, their ability to partition adaptively provides weight- efficient means of balancing gate and value capacity across depth-one that would be impractical with independently parameterized experts.
D PyTorch Implementation of MGLU
PyTorch implementations of the MGLU layer used for training are provided in Algorithm
26

---

1
Percentage (%)
52
50
10111213141511
8 10111213:445116
2 4.5
101112 14151c
Layer
Layer
Layer
mask
mask
mask
mask
mask
mask
mask
Figure 14: Layer-wise gate allocation for learned masks.
Algorithm 2 PyTorch-ttyle Implementation of MGLU (7.m 1).
clasa MCLu Linear):
iniit_ (salf, in_features, aut_features):
super (NGLU, self) initt... (in [in_uuature,, fut_featuuee, False)
self, wet..rrgitte__praaaee
corch.randdcouttfoootte in_featurea), requires_grad-Truu "nask", Paramaatter(O..0
mask binary by straight- through estinator ate maakk soft_nask))
hard_mask mask (soft_ pass > to naek dtype) dtype)
hard_nask (bard_nask soft_mask) detach soft_mack
hard_nask
def forvard(self.. x):
mask self. weweeeeeeeeeelle.oeeee maek)
compute complenentary output: el, c2 el F. F.linear(x, (x, self veight bard_nask)
F.lineer veight return el, c2
hard_nask))
E Efficient Kernel Implementation
straightforward PyTorch implementation of a masked-GLU (MGLU) layer must (i) load the weight matrix from HBM multiple times-once for every mask bit-and (ii) launch several sepa- rate matmul activation operations. Because each extra load traverses the bandwidth-limg limited HBM ++ SRAM link, such code quickly becomes memory-bound. By fusing the entire computation into a single CUDA kernel we touch each weight exactly once, keep partial results in registers, and remove almost all redundant global-memorg traffic.
Implementation Details. Algorithm 3 shows the memory-access access pattern and compute loop for the simplest MGLU case, 7mm = N __MASSS. Each thread block is launched with coordinates (row, chunk) that it processes fetched together as a single half2 load, converted once to float2, and kept in registers for the entire aacliply-accumulate step. The binary masks packed eight bits per int8; testing the active bit yields a 0/1 scalar that is multiplied into the product. Both the ungated sum and each gated sum accumulate only in registers; after the loop a warp-level shuffle reduces these partials and a single atomic write row sends the result to HBM.
output row and one K-slice at time. Two FP16 weights and activations
Head-room Hopper. Although the kernel already eliminates almost all redundant global-memory traffic, it deliberately avoids Hopper-specifi. optimisations such as cp. async and Tensor Memory Accelerator (TMA). Incorporating those features could conservatively deliver additional 1.2-1.3x speed-up on server- server-grade H100 GPUs an avenue we leave to future work.
Latency Comparison. In Tables 16 and 17 we report the execution latency of our kernel versus a naïve Torch implementation on RTX 5090 and H100 GPUs, respectively: Tables T8]and and T9 then compare kernel against the standard GLU implementation the same devices.
27

---

Algorithm 3 Simplified CUDA Implementation of MGLU (nm N_ D_MASKS).
global_. void mw_splitk,_naei nasks, karnal
const helle re.rericct... A, // [M X0 row-najor
h.afff restricct... 3, [X]
const int8_t* restrict... mask, D x N),
flonte .rrestrict.. int H, split_2
naaked outputa
int row blockIdx.X, int chunk blockIdx.y; if (row
chunk split_k) return;
// conpute [start, of this K-chunk
int chunk_aize
split_k 1) split_k;
start chink chunk_size; int minlstart chunk_aize, X0;
int rov_eff = row + chreadIdx.xx 2; stride blockDim.x
idx
flost total 0.0I; flost neun(K MASKS 2] - (0.0f, 0.0E, ...):
for idx end; idx 4- stride)
--half2 *reintarpret, castrrrr ha
idx]);
caskkoontt ..hell2 x2 ..haf22floot2(=2);
(Ax
float2 __bbll(221oot22=22)
flost2 pred af .xxxxxfx x, af.y*xf.y };
total prod.x prod.y;
int8_t ___Idg (rov_ idx] [rou_ot int8_t m1 (roo off idx+i]); unsigned
(unsigned int)m0;
unsigned int bi (unsigned int)mi; #pragna unroll
(int float mb0 float((bO >> 1u); flost mbi float((b1 b) iu); if (b0 (Iu << b)) msun[b] + prod.x; if (bl & (1u << b)) moun [b] prod.y;
N_MASKS; ++b) {
// reduce and write to HBN.
Table 16: MGLU latency and speed-ups, ups. Torch MGLU the naive nn Linear implementation on an RTX 5090 GPU; higher ratios mean faster custom kernels.
T
h
CUDA (ms)
Triton (ms)
Torch (ms)
Torch/CUDA
Torch/Triton
8192
2048
0.0202
0,0516
0,0715
3.54x
1.39
8192
2048
0.0210
0.0533
0.1358
6.47x
2.55
8192
2048
0.0217
0.0606
0.2715
12.51
4.48
8192
2048
0.0265
0.0834
0.5210
19.66
6.25x
14336
4096
0.1166
0.1342
0.3289
2.82x
2.45
14336
4096
0.1169
0.1381
0.6001
5.13x
4.35x
14336
4096
0.1186
0.1454
1.1426
9.63x
7.86x
14336
4096
0.1229
0.1990
2.2172
18.04x
11.14x
Table 17: MGLU latency and speed-ups. Torch MGLU is the naïve nn. Linear implementation on H100 GPU; higher ratios mean faster custom kernels.
h
CUDA (ms)
Triton (ms)
Torch (ms)
Torch/CUDA
Torch/Triton
8192
2048
0.0395
0.0639
0.1296
3.28x
2.03 x
8192
2048
0.0409
0.0679
0.2213
5.41 x
3.26x
8192
2048
0.0428
0.0810
0.4117
9.62x
5,08
8192
2048
0.0483
0.1294
0.7910
16.38
6.11x x
14336
4096
0.1044
0.1191
0.3896
3.73x
3.27 x
14336
4096
0.1067
0.1239
0,7080
6.64x
5,71
14336
4096
0.1110
0.1608
1.3454
12.12x
8.37
14336
4096
0.1215
0.3202
2.6354
21.68x
8.23x
28

---

Table 18: CUDA MGLU latency and speed-ups against the standard PyTorch GLU baseline (no masking) on an RTX 5090 GPU. Lower latency and higher speed-up up are better.
7ym
h
CUDA MGLU (ms)
Torch GLU
GLU/CUDA
8192
2048
0.0202
0.0306
1.51x
8192
2048
0.0210
0.0306
1.46x
8192
2048
0.0217
0.0306
1.41x
8192
2048
0.0265
0.0306
1.15x
14336
4096
0.1166
0.1530
1.31x
14336
4096
0.1169
0.1530
1.31
14336
4096
0.1186
0.1530
1.29
14336
4096
0.1229
0.1530
1.24x
Table 19: CUDA MGLU latency and speed-ups against the standard PyTorch GLU baseline (no masking) H100 GPU. Lower latency and higher speed-up better.
h
CUDA MGLU (ms)
Torch GLU GLU/CUDA
8192
2048
0.0395
0.0485
1.23
8192
2048
0.0409
0.0485
1.19x
8192
2048
0.0428
0.0485
1.13x
8192
2048
0.0483
0.0485
1.00
14336
4096
0.1044
0.1215
1.16x
14336
4096
0.1067
0.1215
1.14x
14336
4096
0.1110
0.1215
1.09 X
14336
4096
0.1215
0.1215
1.00
F Broader Impacts
The method proposed in this paper will lead to negative societal impact. By reducing memory- bandwidth pressure and per-token latency, it cuts the energy consumption of lrgg-langgage--mon inference and therefore lowers carbon emissions. These efficiency gains also make advanced language capabilities viable on commodity and edge hardware, broadening access while trimming operational costs.
29


Reviews:
[{"review_id": "ETGoSPSvRN", "reviewer": "Reviewer_Ukks", "summary": "The authors propose a new activation function / architectural modification of the transformers' feed-forward layer called MGLU (masked gated linear unit). By cleverly implementing element-wise masking, their proposed method shows better performance while being faster than standard activation functions.", "questions": "- What are the wall-clock pretraining times of the experiments in section 5?\n- Do you have any hypotheses as to why MGLU leads to more robust and stable training than GLU?\n- 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?\n- 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?", "limitations": "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.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n- The idea is smart, but what is more important, it is executed well. While a naive implementation of MGLU would be prohibitively slow, the authors paid a lot attention to its efficient kernel implementation, which cleverly uses sparsity and results in even better inference-time efficiency than standard GLU.\n- The paper is written well; it was fun to read and easy to understand.\n- From what I can see in the main part of the paper, the experimental setup is quite robust. Unfortunately, the authors forgot to include the appendices, so it is hard to properly judge the validity of the training and evaluation pipelines,\n\n**Weaknesses**\n- The paper extensively talks about the improved inference-time efficiency, but the training-time computational cost is only briefly mentioned in the paragraph on lines 197–205. If I understand the implementation correctly, MGLU is substantially slower compared than GLU during training because it has to use float16-masks instead of the efficient bit-masks. Is that correct? If so, then that should be clearly highlighted in the paper, talking only about the inference-time efficiency is misleading.", "was_revised": true, "final_justification": "This work presents a well-executed way of increasing inference speed of transformers. The authors clearly responded to all my comments.", "rebuttal": "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.\n\n> The authors forgot to include the appendices, so it is hard to properly judge the validity of the training and evaluation pipelines.\n\nThe 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.\n\n> What are the wall-clock pretraining times of the experiments in section 5?\n\nThe 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.\n\n> Do you have any hypotheses as to why MGLU leads to more robust and stable training than GLU?\n\nAs 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.\n1) MGLU removes high-variance activation fluctuations. This \"smoothing\" of the forward pass reduces gradient noise, allowing more consistent updates and preventing early-training oscillations.\n2) 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.\n\n> 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?\n\nWe'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.\n\n| Model | $n_m$ | ArcE | ArcC | Hellaswag | PIQA | SCIQ | Winogrande | Avg |\n| -- | -- | -- |  -- |  -- |  -- |  -- |  -- |  -- |\n| SwiMGLU (independent masks) | 4 | 65.32 | 32.42 | 36.32 | 69.31 | 88.70 | 52.64 | 57.45 |\n| SwiMGLU (complementary masks) | 4 | 67.13 | 30.55 | 37.53 | 69.48 | 88.60 | 53.91 | 57.87 |\n\n> 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?\n\nWhile 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.\n\nThe 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.\n\nBased 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.\n\n| Model | $h$ | $d$ | $n_m$ | train/loss | val/ppl | avg score | \n| ---- | ---- | ---- | -------- | ------------ | -------- | -------- |\n| SwiGLU    | 14336 | 4096 | - | 3.17 | 23.4 | 40.50 |\n| SwiMGLU | 14336 | 4096 | 4 | 3.11 | 22.2 | 42.02 |\n\n> 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.\n\nThank 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."}, {"review_id": "t27NRvtXX3", "reviewer": "Reviewer_FrzZ", "summary": "This paper proposes an interesting implementation, Masked gated linear units, to replace the SwiGLU, with the aim of sharing the projection and gate layer for storage reduction and latency reduction. The authors keep the efficacy of the approximation by introducing multiple branches and further address efficiency concerns with a customized CUDA implementation to share one weight matrix multiplication with indexing ops.", "questions": "Questions\n\n* What would be the overall runtime breakdown, and how much gain can we get by switching to the proposed implementation instead of SwiGLU?\n* It seems the performance still lags behind the SwiGLU, which I may be more convinced of when seeing larger models’ results.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strength**\n* I appreciate the efforts to replace the basic operators in LLM, which is often treated as a golden paradigm.\n* The author tries to approximate the SwiGLU with comparable performance and inference speed benefits.\n\n**Weakness**\n* One concern is that the introduction of hyperparameters, such as branch number, may make it tricky to choose, especially if there seems to be no fixed rule to pick the best n, as shown in the Table.3\n* Another concern is the limited validation of the proposed method, which is conducted on a small-scale model. More model series and larger model scales should be conducted.\n* My last concern is that there seems to be no theoretical connection as to why sharing the same weights by a mask can obtain similar performance as SwiGLU. It will make the design unintuitive and hard to understand or adopt.", "was_revised": true, "final_justification": "It addresses my concern, and I will give a positive rating.", "rebuttal": "We sincerely appreciate your valuable feedback.\n\n> What would be the overall runtime breakdown, and how much gain can we get?\n\nThe 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.\n\n> It seems the performance still lags behind the SwiGLU, I may be more convinced when seeing larger models’ results.\n\nPre-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.\n\n> Hyperparameter (branch number) may be tricky to choose.\n\nBased 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.\n\n> Limited validation of the proposed method, which is conducted on a small-scale model. More model series and larger model scales should be conducted.\n\nResults using GeGLU (as in the Gemma series) are provided in the Appendix, demonstrating that our method remains consistently effective. Scalability is discussed above.\n\n> There seems to be no theoretical connection as to why sharing the same weights by a mask can obtain similar performance as SwiGLU.\n\nSharing 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.\n\n[1] E. Malach, et al. Shamir, Proving the Lottery Ticket Hypothesis: Pruning is All You Need, ICML, 2020.\n\nWe hope our responses have adequately addressed the reviewer's concerns."}, {"review_id": "Mv0HoVXnHB", "reviewer": "Reviewer_NqHP", "summary": "This paper introduces a new variant of GLU focusing on efficiency. It features a mixture of element-wise gating architecture and a hardware-friendly kernel. Experiments indicate that the proposed approach achieves efficiency improvements while maintaining good quality.", "questions": "See earlier.", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "In terms of strength, this paper is well-written, with a well-motivated algorithm followed by well-executed experiments, both in terms of performance and efficiency. The method is also co-designed with kernel implementation in mind, which is commendable.\n\nRegarding areas for improvement, I think the positioning of the paper could be clearer. It took me some time to realize that the main advantages over existing architectures are that it reduces parameter counts without sacrificing quality, essentially creating a sparse-by-design architecture. Notably, speed advantages over existing architectures aren't its main selling point (which is fine with me). The introduction and abstract consistently emphasize efficiency compared to a naive implementation of its own architecture. I initially misunderstood it as a substantial (>10x) speedup over existing (Swi)GLU.\n\nQuestions:\n\nRegarding FlashMGLU:\n- Line 163 states that the IO is $n_m \\times 2 + 1$, but it becomes $n_m \\times 4$. Could the author clarify this?\n- 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?\n- 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.\n- Line 254, what is the speedup of `1.51x` compared against? (Is it standard SwiGLU or GLU?)\n- 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`?\n- Line 194: Is it `3hd` or `2hd`?\n- Line 195: Is it `2hd` or `16hd`?\n\nMore questions:\n- Are the masks input-dependent or independent (it seems to be the latter)?\n- 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.) \n- 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)?", "was_revised": true, "final_justification": "This work proposes a well-motivated method with great execution. The authors also clarify a few places during the rebuttal.", "rebuttal": "We appreciate the reviewer's careful reading and detailed feedback. Below, we respond to each question.\n\n> Line 163 states that the IO is $n_m \\times 2 + 1$, but it becomes $n_m \\times 4$. Could the author clarify this?\n\nThe 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$.\n\n> 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?\n\nFrom 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.\n\n> 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.\n\nOur 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.\n\n> Line 254, what is the speedup of 1.51x compared against? (Is it standard SwiGLU or GLU?)\n\nThis 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.\n\n> 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?\n\nThe 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.\n\n> Line 194: Is it 3hd or 2hd?\n> Line 195: Is it 2hd or 16hd?\n\nWe 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.\n\n### Additional questions:\n\n> Are the masks input-dependent or independent (it seems to be the latter)?\n\nThe 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.\n\n> 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.)\n\nIn 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.\n\n> 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)?\n\nWhile 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."}, {"review_id": "rsKLHb2eBQ", "reviewer": "Reviewer_GUVW", "summary": "This paper proposes Masked Gated Linear Units (MGLUs), a novel feed-forward layer design aimed at reducing inference-time memory usage and improving computational efficiency in large language models (LLMs), without sacrificing performance. MGLUs share a single matrix for both gating and projection operations, while using a masking mechanism to allocate weights for different purposes.", "questions": "See weakness.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strength:**\n\n1. The paper is well written and easy to read.\n\n2. The proposed method replaces two projections with a single weight matrix, reducing memory bandwidth demand.\n\n3. The extensive evaluation show that MGLU can achieve comparable performance to SwiGLU but with less memory loading.\n\n\n**Weakness and Questions:**\n\n1. 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.\n\n2. What is the end-to-end model inference speed-up?", "was_revised": true, "final_justification": "The response has addressed my concerns.", "rebuttal": "We thank the reviewer for acknowledging the effectiveness of our MGLU layer. We respond to the reviewer’s questions below.\n\n> 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.\n\nTo 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.\nEven against an optimized MGLU baseline, FlashMGLU still yields up to 12.23× speedup in the CUDA kernel.\n\n### RTX 5090\n| $h$    | $d$    | $n_m$ | CUDA (ms) | Triton (ms) | Torch (ms) | Torch/CUDA | Torch/Triton |\n|------|------|------|------------|--------------|-------------|--------------|------------------|\n| 8192 | 2048 | 1    | **0.0202** | 0.0516       | 0.0837      | **4.14 $\\times$** | 1.62 $\\times$ |\n| 8192 | 2048 | 2    | **0.0210** | 0.0533       | 0.1337      | **6.37 $\\times$** | 2.51 $\\times$ |\n| 8192 | 2048 | 4    | **0.0217** | 0.0606       | 0.1933      | **8.91 $\\times$** | 3.19 $\\times$ |\n| 8192 | 2048 | 8    | **0.0265** | 0.0834       | 0.3240      | **12.23 $\\times$** | 3.88 $\\times$ |\n|14336 | 4096 | 1    | **0.1166** | 0.1342       | 0.3554      | **3.05 $\\times$** | 2.65 $\\times$ |\n|14336 | 4096 | 2    | **0.1169** | 0.1381       | 0.4660      | **3.99 $\\times$** | 3.37 $\\times$ |\n|14336 | 4096 | 4    | **0.1186** | 0.1454       | 0.6906      | **5.82 $\\times$** | 4.75 $\\times$ |\n|14336 | 4096 | 8    | **0.1229** | 0.1990       | 1.1417      | **9.29 $\\times$** | 5.74 $\\times$ |\n\n### H100\n| $h$    | $d$    | $n_m$ | CUDA (ms) | Triton (ms) | Torch (ms) | Torch/CUDA | Torch/Triton |\n|------|------|------|------------|--------------|-------------|-------------|----------------|\n| 8192 | 2048 | 1    | **0.0395** | 0.0639       | 0.1321      | **3.34 $\\times$** | 2.07 $\\times$ |\n| 8192 | 2048 | 2    | **0.0409** | 0.0679       | 0.1639      | **4.01 $\\times$** | 2.41 $\\times$ |\n| 8192 | 2048 | 4    | **0.0428** | 0.0810       | 0.2215      | **5.18 $\\times$** | 2.73 $\\times$ |\n| 8192 | 2048 | 8    | **0.0483** | 0.1294       | 0.3260      | **6.75 $\\times$** | 2.52 $\\times$ |\n|14336 | 4096 | 1    | **0.1044** | 0.1191       | 0.3399      | **3.26 $\\times$** | 2.85 $\\times$ |\n|14336 | 4096 | 2    | **0.1067** | 0.1239       | 0.4212      | **3.95 $\\times$** | 3.40 $\\times$ |\n|14336 | 4096 | 4    | **0.1110** | 0.1608       | 0.6125      | **5.52 $\\times$** | 3.81 $\\times$ |\n|14336 | 4096 | 8    | **0.1215** | 0.3202       | 0.9882      | **8.13 $\\times$** | 3.09 $\\times$ |\n\n> What is the end-to-end model inference speed-up?\n\nWith 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.\n\n| Model    |   $h$   |   $d$   | $n_m$ | Throughput (token/sec) | Speed-up vs. SwiGLU |\n|----------|--------|--------|--------|--------------------------|--------------------------|\n| SwiGLU   |  8192   |  2048   | –      | 396.9                    | 1.00 $\\times$            |\n| SwiMGLU  |  8192   |  2048   | 1      | 445.9                    | 1.12 $\\times$            |\n| SwiMGLU  |  8192   |  2048   | 2      | 446.1                    | 1.12 $\\times$            |\n| SwiMGLU  |  8192   |  2048   | 4      | 441.5                    | 1.11 $\\times$            |\n| SwiMGLU  |  8192   |  2048   | 8      | 442.8                    | 1.12 $\\times$            |"}]

Author rebuttal:

Meta-review: This paper proposes Masked GLU (MGLU), a replacement for the popular gated linear unit FFN (SwiGLU) that is now de facto standard in Transformers. Unlike SwiGLU which maintains separate projection matrices for the gating component, MGLU uses a learned binary mask on the same matrix to obtain the gating values. This enables parameter-efficiency, as well as inference-time speedups (due to less memory movement). The approach is tested and found to work well on reasonable-scale language modeling experiments.

On the plus side, this is a simple but clever idea that roughly maintains the performance of SwiGLU while being more parameter- and memory-efficient. The empirical experiments are thorough, the the reviewers (and the AC) appreciated the kernel which results in real gains over strong SwiGLU baselines.

On the negative side, MGLU is significantly more expensive during training time, which might prevent its application in real-world language models, where training compute is a significant consideration. (Nonetheless, there are many scenarios where training compute is not really a factor, e.g., in training smaller LLMs for the edge, where this method could really shine). Moreover, while not mentioned by reviewers, the AC notes that it is possible that this type of weight sharing could make post-training sparsification/quantization more difficult---I encourage the authors to investigate this aspect.

Overall, this is a very well-executed paper that should be of high interest to the architecture community. A clear accept.
Decision: Accept (poster)