龚卓成关于transformer模型架构的论文《Mixture-of-Modules: Reinventing Transformers as Dynamic Assemblies of Modules》被EMNLP 2024接收。目前的Transformer模型中从浅层到深层按固定顺序调用所有模型层。在这项工作中，我们尝试打破深度有序的惯例，提出了一种称为混合模块（Mixture-of-Modules，MoM）的新架构。MoM的构建从由多头注意力和前馈网络定义的有限模块集开始。然后，两个路由模块迭代地从该集合中选择注意力模块和前馈模块来处理token。选择过程动态地扩展了token前向传播中的计算图，最终形成模块的组合。MoM不仅为Transformer及其众多变体提供了统一的框架，还提供了一种灵活且可学习的方法。我们使用OpenWebText对各种MoM进行预训练。经验结果表明，不同大小的MoM始终优于vanilla transformer。更有趣的是，即使移除50%的多头注意力模块和25%的前馈模块，MoM模型仍保持可比的性能。此外，通过适当调整模块数量和压缩模型深度，可以得到一个与GPT-2（774M）具有可比性能的MoM，同时在前向计算期间节省16%的TFLOPs和42%的内存使用。