掩码门控线性单元
机器学习
2025-07-01 v1 计算与语言
摘要
门控线性单元（GLUs）已成为最先进的大语言模型（LLMs）前馈网络中的关键组件。然而，由于门控流与值流使用独立的权重矩阵，与无门控的前馈层相比，它们需要两倍的内存读取次数。为解决这一瓶颈，我们提出了掩码门控线性单元（MGLUs），这是一种具有高效内核实现的新型 GLU 家族。MGLUs 的核心贡献包括： 元素级门控混合架构，该架构在单一共享权重矩阵上学习多个二值掩码，每个掩码在元素级别决定门控或值的分配，从而减少内存传输； FlashMGLU，这是一种硬件友好的内核，相较于朴素的 PyTorch MGLU 实现可实现高达 19.7 的推理时加速，并且在 RTX5090 GPU 上，尽管增加了架构复杂度，仍比标准 GLUs 节省 47% 的内存且速度提升 34%。在 LLM 实验中，Swish 激活的变体 SwiMGLU 在保持内存优势的同时，其下游准确率与 SwiGLU 基线相当——甚至有所超越。
引用
@article{arxiv.2506.23225,
title = {Masked Gated Linear Unit},
author = {Yukito Tajima and Nakamasa Inoue and Yusuke Sekikawa and Ikuro Sato and Rio Yokota},
journal= {arXiv preprint arXiv:2506.23225},
year = {2025}
}