Research map / Architectures and efficiency / RNNs and normalization
Masked Gated Linear Unit
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Masked Gated Linear Units (MGLUs), which replace the two GLU projection matrices with one shared weight matrix plus learnable binary masks (MoEG), and a fused FlashMGLU kernel that cuts inference memory reads.
Problem
GLU variants used in LLM feed-forward layers require two full-precision weight matrices, doubling memory reads during memory-bound inference.
Approach
Learn n_m binary masks over a single shared weight matrix W so that each mask splits W into complementary gate and value subspaces (Mixture of Element-wise Gating), and fuse mask unpacking with the matrix-vector product in a Split-K kernel (FlashMGLU) so each weight is read once.
Key results
- SwiMGLU with 4 masks reaches 56.85% zero-shot average vs. SwiGLU 56.00% at the 1.08B scale (808M weights).
- FlashMGLU achieves up to 19.66x speed-up over naive PyTorch MGLU on RTX 5090 (n_m=8, large setting), and is reported 34% faster and 47% more memory-efficient than standard GLUs.
- At the small 159M scale, SwiMGLU with n_m=4 reaches 23.9 PPL vs. SwiGLU 23.7 PPL with 113M vs. 141M weights.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace SwiGLU with SwiMGLU using 4 learned masks | Llama 3 large · Downstream zero-shot | Avg accuracy (%): 56.00 -> 56.85 | Helped |
| Replace SwiGLU with SwiMGLU using 4 learned masks | Llama 3 large · Downstream two-shot | Avg accuracy (%): 57.36 -> 57.87 | Helped |
| Increase the learned mask count from 1 to 8 (ablation) | SwiMGLU, Llama 3 small · Downstream two-shot | Avg accuracy (%): 44.71% to 46.74% | Helped |
| Increase the learned mask count from 1 to 8 (ablation) | SwiMGLU, Llama 3 small · Validation perplexity | Perplexity: 25.0 -> 23.5 | Helped |
| Learn masks instead of fixing them randomly, with 2 masks (ablation) | SwiMGLU, Llama 3 small · Downstream zero-shot | Avg accuracy (%): 45.84 -> 46.62 | Helped |
| Learn masks instead of fixing them randomly, with 2 masks (ablation) | SwiMGLU, Llama 3 small · Validation perplexity | Perplexity: 25.1 -> 24.5 | Helped |
| Use fixed rather than learned masks, with 1 mask (ablation) | SwiMGLU, Llama 3 small · Downstream zero-shot | Avg accuracy (%): 46.06 -> 46.30 | Mixed |
| Route each token to only 1 of 4 masks instead of summing all 4 (ablation) | SwiMGLU, Llama 3 small · Downstream zero-shot | Avg accuracy (%): 46.48 -> 46.07 | Hurt |
| Route each token to 2 of 4 masks instead of summing all 4 (ablation) | SwiMGLU, Llama 3 small · Downstream zero-shot | Avg accuracy (%): 46.48 -> 46.64 | Helped |
| Use GELU rather than Swish as the MGLU gating activation (4 masks) (ablation) | SwiMGLU, Llama 3 small · Downstream zero-shot | Avg accuracy (%): 46.48 -> 46.74 | Helped |
| Use ReLU rather than Swish as the MGLU gating activation (4 masks) (ablation) | SwiMGLU, Llama 3 small · Downstream zero-shot | Avg accuracy (%): 46.20 -> 46.38 | No clear effect |
| Remove the gate mask (dense gate) (ablation) | Single-mask MGLU, Llama 3 small · Training and validation | Loss and perplexity | Hurt |
Surprising findings
- Naively sharing one weight matrix across SwiGLU gate and value streams raises perplexity from 23.6 to 27.0, so explicit gating via masks is needed.
- Freezing the random masks collapses learning to the single-mask baseline, suggesting learned mask combinations rather than sparsity drive the gains.
Limitations
- MGLU increases arithmetic complexity and training cost (masks stored as FP16 logits), with training overhead growing with mask count.
- Performance advantage over standard GLU kernel erodes once n_m=8 causes register spilling; Triton prototype is slower on small matrices.
- Downstream gains are small and noisy on small models; SwiMGLU sometimes underperforms SwiGLU (e.g., n_m=1, 2 at 1.08B scale).
Experimental setup
Datasets: FineWeb-Edu 100B
Benchmarks: ARC Easy ARC Challenge HellaSwag PiQA SciQ Winogrande
Metrics: zero-shot accuracy two-shot accuracy validation perplexity latency memory load
Baselines: GELU FFN SwiGLU naive PyTorch MGLU standard GLU nn.Linear kernel Triton prototype
Builds on: GLU variants SwiGLU straight-through estimator mixture of experts Switch Transformer
Closest papers by meaning
- Liger: Linearizing Large Language Models to Gated Recurrent Structures (ICML 2025 · similarity 0.87)
- Gated Linear Attention Transformers with Hardware-Efficient Training (ICML 2024 · similarity 0.86)
- MaskPro: Linear-Space Probabilistic Learning for Strict (N:M)-Sparsity on LLMs (ICLR 2026 · similarity 0.86)
- GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection (ICML 2024 · similarity 0.86)
- Search for Efficient Large Language Models (NeurIPS 2024 · similarity 0.85)
- Learn To be Efficient: Build Structured Sparsity in Large Language Models (NeurIPS 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.