Gemma4-Pruning：基于互信息与Rényi熵的无监督大语言模型结构化剪枝方案

一个无需重新训练的无监督结构化剪枝管道，通过互信息估计和基于矩阵的Rényi熵计算，实现对Gemma 4模型的FFN层压缩，在保持模型性能的同时显著降低推理成本。

• 原作者/维护者：jaja07
• 来源平台：github
• 原始标题：Gemma4-Pruning
• 原始链接：https://github.com/jaja07/Gemma4-Pruning
• 来源发布时间/更新时间：2026-06-10T16:13:04Z

原作者与来源

• 原作者/维护者：jaja07
• 来源平台：github
• 原始标题：Gemma4-Pruning
• 原始链接：https://github.com/jaja07/Gemma4-Pruning
• 来源发布时间/更新时间：2026-06-10T16:13:04Z 原作者与来源\n\n- 原作者/维护者：jaja07\n- 来源平台：GitHub\n- 原始标题：Gemma4-Pruning\n- 原始链接：https://github.com/jaja07/Gemma4-Pruning\n- 发布时间：2026年6月10日\n\n背景与动机\n\n随着大语言模型（LLM）规模不断膨胀，如何在保持模型能力的前提下降低推理成本成为关键挑战。传统的模型压缩方法往往需要昂贵的重新训练流程，而结构化剪枝作为一种"一次性"压缩技术，能够在不重新训练的情况下直接移除冗余神经元，大幅降低显存占用和计算开销。\n\nGemma4-Pruning项目针对Google的Gemma 4模型（4B参数版本），提出了一种完全无监督的剪枝方案。该方法无需标注数据，仅需少量未标注样本即可评估神经元间的信息冗余度，从而实现高效的模型压缩。\n\n核心方法论\n\n基于信息论的神经元冗余评估\n\n该项目的核心创新在于使用互信息（Mutual Information）来量化神经元之间的功能重叠。传统剪枝方法通常依赖权重大小或梯度信息，而互信息能够直接捕捉神经元激活模式之间的统计依赖性，更准确地识别功能相似的冗余神经元。\n\n项目采用基于再生核希尔伯特空间（RKHS）的互信息估计器，结合矩阵形式的Rényi熵计算，实现了对高维神经元激活分布的有效建模。\n\nRényi熵的矩阵形式定义\n\n对于观察到的神经元激活向量$Z_k$，项目首先构建高斯RBF核的Gram矩阵：\n\n$$K_k(i, j) = \exp\left(-\frac{(z_{k,i} - z_{k,j})^2}{2\sigma^2}\right)$$\n\n通过迹归一化得到$\tilde{K}k = \frac{K_k}{\mathrm{tr}(K_k)}$后，Rényi熵的计算转化为对归一化Gram矩阵特征值的幂运算：\n\n$$S\alpha(\tilde{K}k) = \frac{1}{1-\alpha} \log_2\left(\sum{i=1}^{N} \lambda_i^\alpha\right)$$\n\n这种方法避免了传统直方图估计在高维空间中的维度灾难问题，同时保持了熵的数学性质。\n\n互信息的核化估计\n\n两个神经元之间的互信息通过联合熵计算：\n\n$$I(Z_k ; Z_l) = S_\alpha(\tilde{K}k) + S\alpha(\tilde{K}l) - S\alpha(\tilde{K}{\text{joint}})$$\n\n其中联合核矩阵通过Hadamard积构造：$K{\text{joint}} = K_k \odot K_l$。这种构造方式保证了联合分布的核函数仍然是正定核，满足RKHS理论的要求。\n\n自适应核宽度与聚类剪枝\n\n项目采用Scott规则自适应确定核宽度：\n\n$$\sigma = \gamma N^{-\frac{1}{4+d}}$$\n\n其中$N$为样本数，$d$为维度，$\gamma$为可调超参数。这种自适应机制使得方法能够适应不同规模的激活数据。\n\n在计算完所有神经元对的互信息后，项目使用多维缩放（MDS）将高维距离映射到低维空间，然后应用KMeans聚类。在每个聚类簇中，仅保留最接近质心的代表神经元，其余神经元则从FFN层的gate_proj、up_proj和down_proj权重矩阵中物理移除。\n\n实现架构\n\n项目采用模块化设计，主要包含以下组件：\n\n数据准备模块\n\nscripts/download_dataset.py负责从Hugging Face下载ECE-ILAB/resilient-ai-unified数据集，并使用datasets.save_to_disk保存为本地格式。这一设计允许用户在无网络连接的环境下进行剪枝操作。\n\n剪枝执行模块\n\nscripts/prune.py是核心执行脚本，其工作流程包括：\n\n1. 模型加载：使用Transformers库加载google/gemma-4-E4B-it模型，默认以bfloat16精度加载以减少显存占用\n2. 激活捕获：在FFN层的down_proj上注册前向钩子，捕获神经元激活值\n3. 子采样：对token序列进行子采样，控制内存使用量\n4. 互信息计算：分块计算神经元间的互信息矩阵\n5. 聚类与剪枝：基于互信息进行聚类，识别并移除冗余神经元\n6. 模型导出：保存剪枝后的模型，同时更新Hugging Face配置中的intermediate_size参数\n\n配置管理\n\n项目使用pyproject.toml管理依赖，支持通过uv工具进行环境管理。CUDA 12.4的PyTorch索引已预先配置，简化了GPU环境的搭建流程。\n\n技术亮点与创新点\n\n无需重新训练\n\n与许多剪枝方法不同，Gemma4-Pruning在剪枝后无需任何微调或重新训练即可使用。这得益于互信息估计的准确性——通过识别功能真正相似的神经元进行剪枝，而非简单基于权重大小，使得剩余神经元已经具备了足够的表达能力。\n\n纯无监督设计\n\n方法完全不依赖标注数据，仅需少量未标注文本即可估计神经元激活分布。这使得剪枝过程可以应用于任何领域，无需准备特定任务的标注数据。\n\n数学严谨性\n\n项目基于信息论的严格数学框架，使用RKHS理论和Rényi熵的矩阵形式，避免了传统方法中的启发式假设。这种理论基础保证了方法在不同模型和数据分布下的稳定性。\n\n物理剪枝而非掩码\n\n与许多仅添加稀疏掩码的"软剪枝"方法不同，Gemma4-Pruning真正从权重矩阵中移除冗余行/列，实现模型尺寸的实质性缩减。这不仅减少了存储空间，还能带来实际的推理加速。\n\n使用方式与注意事项\n\n项目使用简单，只需两步即可完成剪枝：\n\npowershell\nuv run python scripts/download_dataset.py\nuv run python scripts/prune.py\n\n\n默认配置使用google/gemma-4-E4B-it模型，输出路径为models/gemma4-pruned-mi/。用户可根据需要修改脚本中的路径配置。\n\n注意事项：\n\n- 剪枝过程需要较大显存，项目已配置device_map=\"auto\"以支持多GPU自动分配\n- 依赖库包括torch、transformers、datasets、scikit-learn、numpy和tqdm\n- 当前实现主要针对Gemma 4的FFN层，未来可扩展至注意力头剪枝\n\n实践意义与展望\n\nGemma4-Pruning为大语言模型的边缘部署提供了可行路径。通过将4B参数模型压缩至更小尺寸，该技术使得在资源受限设备上运行高质量LLM成为可能。\n\n该项目的互信息估计框架具有通用性，可扩展至其他模型架构（如Llama、Mistral等）和不同剪枝粒度（如注意力头、层级别）。未来发展方向可能包括：\n\n- 支持更大规模模型的剪枝\n- 引入任务特定的微调阶段以恢复剪枝损失\n- 探索动态剪枝策略，根据输入自适应激活不同子网络\n\n对于希望降低LLM部署成本的研究者和开发者而言，Gemma4-Pruning提供了一个即插即用的开源解决方案。