门控注意力：实现与可视化

本仓库包含基于 Qwen3 模型架构的 门控注意力 机制的实现，以及用于可视化注意力图的工具。我们的修改基于近期研究发现：在缩放点积注意力（SDPA）之后应用 稀疏的、头级别的门控机制，能够显著提升性能、训练稳定性和长上下文泛化能力。更多细节请参阅我们的论文《面向大型语言模型的门控注意力：非线性、稀疏性与无注意力陷阱》（arXiv:2505.06708）。

---

🆕 最新动态

2025年12月20日 — 我们将发布关于门控注意力有效性的更多分析和案例研究，包括可视化结果及对模型内部机制的定量探究，以更直观地揭示门控机制的作用原理及其有效性。

2025年11月26日 — 我们的论文 面向大型语言模型的门控注意力：非线性、稀疏性与无注意力陷阱 荣获 NeurIPS 2025 最佳论文奖！🎉
该奖项从5,290篇被接受的论文中仅评选出4篇获奖，彰显了我们在注意力机制设计方面的基础性贡献。官方公告：NeurIPS 2025 最佳论文奖

2025年9月18日 — 我们的论文 面向大型语言模型的门控注意力：非线性、稀疏性与无注意力陷阱 被选为 NeurIPS 2025 的口头报告，入选比例仅为1.5%（在5,290篇被接受的论文中排名第77位）。这一荣誉充分肯定了我们关于重新思考注意力门控机制以构建可扩展、稳定且具备长上下文处理能力的语言模型的研究成果的重要性与创新性。

2025年9月10日 — 门控注意力 已成功集成到官方的 Qwen3-Next 架构中，并在 Qwen 的最新研究博客中有所介绍（阅读此处），同时部署于 Qwen3-Next-80B-A3B-Instruct 模型中。这一实际应用验证了我们的核心假设：门控机制能够显著提升 训练稳定性 和 超长上下文性能（最高可达100万 tokens）。

---

📚 引言

门控机制长期以来一直是神经网络设计中的重要组成部分，它能够动态控制信息流动。在本工作中，我们专注于将此类机制整合并评估于 Transformer 模型的标准 softmax 注意力层中。

我们提出了一种在 SDPA 输出后引入的 查询依赖型稀疏门控（记为 G1），该门控通过 sigmoid 函数独立调节每个注意力头。这一简单而有效的改动：

• 为由值投影和输出投影形成的低秩变换引入了 非线性；
• 实现了 输入依赖的稀疏性，从而避免了“注意力陷阱”现象——即早期 token 过度主导注意力分布；
• 提升了 训练稳定性，允许使用更大的学习率；
• 增强了 长上下文外推能力，在 RULER 等基准测试中表现出显著优势。

---

📦 模型

我们提供了遵循 Qwen3 架构的不同门控配置模型：

• baseline：标准注意力，不施加任何门控；
• gate_headwise：在 SDPA 后应用头级别门控；
• gate_elementwise：在 SDPA 后应用元素级别门控。

这些模型已在 Hugging Face 仓库 上发布。

---

🧪 示例用法

我们提供了一个示例脚本，用于加载训练好的模型并可视化启用门控后的注意力图。

需求

pip install transformers matplotlib numpy torch

运行示例

python demo.py

运行后将生成名为 {model_name}_selected_layer_attention_maps.png 的文件，展示四个关键层的注意力图。

注意力图对比

以下是三种不同模型变体——baseline、gate_headwise 和 gate_elementwise——在第1层、第7层、第21层和第28层的注意力图。这些可视化有助于说明门控机制如何影响注意力模式，尤其是在“注意力陷阱”现象方面。

在 baseline 模型中，我们观察到明显的“注意力陷阱”效应——第一个 token 在多层中始终获得不成比例的高注意力分数。这表明模型过度依赖初始 token，可能限制了其在其他位置上合理分配注意力的能力。

基线模型

观察：强烈的对角线 dominance，显著关注第一个 token（注意力陷阱）。这种模式在多层中持续存在。

头级别门控模型

观察：头级别门控降低了注意力陷阱效应。注意力分布更加均匀，且更具上下文相关性。

元素级别门控模型

观察：元素级别门控进一步增强了注意力模式的稀疏性和选择性，使得注意力图更为清晰、结构化。

---

📁 仓库结构

qwen3-gated/
├── figs/                    # 论文中使用的图片
├── modeling_qwen3.py        # 修改后的带门控注意力的 Qwen3 模型
├── configuration_qwen3.py   # 包含门控开关的模型配置
├── demo.py                  # 用于加载模型和提取注意力图的简单示例
└── README.md                # 您当前所在的位置

---

🔧 实现细节

实现门控注意力的核心改动位于提供的代码中的 Qwen3Attention 类中。

🧠 门控变体

我们支持两种主要的门控类型：

1. 头级别门控

每个注意力头拥有独立的门控标量。

self.headwise_attn_output_gate = True

这些选项可在模型配置中进行设置：

{
  "headwise_attn_output_gate": true,
  "elementwise_attn_output_gate": false
}

2. 元素级别门控

注意力输出中的每个元素被独立调节。

self.elementwise_attn_output_gate = True

这些选项可在模型配置中进行设置：

{
  "headwise_attn_output_gate": false,
  "elementwise_attn_output_gate": true
}

---

📝 引用

如果您在研究中使用了本代码或模型，请引用我们的论文：

@misc{qiu2025gatedattentionlargelanguage,
      title={Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free}, 
      author={Zihan Qiu and Zekun Wang and Bo Zheng and Zeyu Huang and Kaiyue Wen and Songlin Yang and Rui Men and Le Yu and Fei Huang and Suozhi Huang and Dayiheng Liu and Jingren Zhou and Junyang Lin},
      year={2025},
      eprint={2505.06708},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2505.06708}, 
}

📬 联系方式

如有任何问题或合作意向，欢迎随时通过 qzh11628@gmail.com 与我们联系。

Gated Attention 快速上手指南

本指南帮助开发者快速部署并体验基于 Qwen3 架构的 Gated Attention（门控注意力）机制。该机制通过在缩放点积注意力（SDPA）后引入稀疏的门控操作，有效解决了“注意力汇聚（Attention Sink）”问题，显著提升了模型的训练稳定性及长上下文处理能力。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux / macOS / Windows
• Python 版本：推荐 Python 3.8 及以上
• 硬件要求：建议使用支持 CUDA 的 NVIDIA GPU 以获得最佳推理和可视化性能（CPU 亦可运行演示脚本，但速度较慢）
• 前置依赖：需安装 transformers, torch, matplotlib, numpy

💡 国内加速建议 推荐使用清华或阿里镜像源安装依赖，以提升下载速度：

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple transformers matplotlib numpy torch

若拉取 Hugging Face 模型受阻，建议设置环境变量使用国内镜像：

export HF_ENDPOINT=https://hf-mirror.com

安装步骤

本项目无需复杂的编译过程，只需安装必要的 Python 库并克隆代码库即可。

1. 克隆仓库

   git clone https://github.com/QwQZh/gated_attention.git
   cd gated_attention
   

2. 安装依赖 执行以下命令安装运行演示脚本所需的库：

   pip install transformers matplotlib numpy torch
   

3. 获取模型权重 演示脚本默认会尝试从 Hugging Face 加载预训练模型。确保网络通畅或已配置国内镜像加速。

   • 模型仓库地址：https://huggingface.co/QwQZh/gated_attention
   • 可选变体：baseline (基线), gate_headwise (头级门控), gate_elementwise (元素级门控)

基本使用

项目提供了一个简单的演示脚本 demo.py，用于加载模型并生成注意力热力图，直观对比不同门控机制的效果。

运行演示

在项目根目录下直接运行：

python demo.py

结果查看

脚本运行完成后，将在当前目录生成名为 {model_name}_selected_layer_attention_maps.png 的图片文件。

• 图片内容：展示模型第 1、7、21、28 层的注意力分布图。
• 效果对比：
  • Baseline：可见明显的“注意力汇聚”现象（首 token 占据过高权重）。
  • Gate Headwise/Elementwise：注意力分布更加均匀且具备上下文依赖性，有效消除了首 token 的异常主导，体现了门控机制带来的稀疏性和选择性优势。

自定义配置（进阶）

如需修改门控类型，可调整 configuration_qwen3.py 中的配置参数，或在加载模型时传入相应参数：

• 启用头级门控：

  {
    "headwise_attn_output_gate": true,
    "elementwise_attn_output_gate": false
  }
  

• 启用元素级门控：

  {
    "headwise_attn_output_gate": false,
    "elementwise_attn_output_gate": true
  }