基于互信息与Rényi熵的Gemma 4大模型无监督结构化剪枝

章节 01

导读 / 主楼：基于互信息与Rényi熵的Gemma 4大模型无监督结构化剪枝

原作者与来源

原作者/维护者: jaja07
来源平台: GitHub
原始标题: Gemma4-Pruning
原始链接: https://github.com/jaja07/Gemma4-Pruning
发布时间: 2026年6月10日

背景：大模型压缩的迫切需求

随着大语言模型（LLM）规模不断膨胀，如何在保持性能的前提下降低推理成本已成为AI工程领域的核心挑战。Gemma 4作为Google发布的40亿参数开源模型，虽然性能优异，但其部署仍需要相当的计算资源。

模型剪枝（Pruning）作为一种经典的模型压缩技术，传统方法往往需要昂贵的重训练过程，或者依赖有标签数据来评估剪枝后的性能损失。本项目提出了一种全新的无监督、免重训练结构化剪枝方案，通过信息论方法识别并移除FFN（前馈网络）层中的冗余神经元。

核心创新：互信息与Rényi熵的融合

该项目的算法核心是基于再生核希尔伯特空间（RKHS）中定义的互信息来无监督估计FFN神经元之间的冗余度。

核方法与Gram矩阵构建

对于在$N$个样本上观察到的神经元$Z_k$，首先使用RBF高斯核构建Gram矩阵：

[K_k(i, j) = \exp\left(-\frac{(z_{k,i} - z_{k,j})^2}{2\sigma^2}\right)]

经过迹归一化后，$\tilde{K}_k = \frac{K_k}{\mathrm{tr}(K_k)}$，可以计算Rényi熵：

[S_\alpha(\tilde{K}k) = \frac{1}{1-\alpha} \log_2\left(\sum{i=1}^{N} \lambda_i^\alpha\right)]

其中${\lambda_1, \dots, \lambda_N}$是$\tilde{K}_k$的特征值。

互信息计算与冗余识别

两个神经元$k$和$l$的联合熵通过其核矩阵的Hadamard积估计：$K_{\text{joint}} = K_k \odot K_l$。互信息则为：

[I(Z_k ; Z_l) = S_\alpha(\tilde{K}k) + S\alpha(\tilde{K}l) - S\alpha(\tilde{K}_{\text{joint}})]

核宽度$\sigma$遵循基于样本数$N$和维度$d$的Scott经验规则：

[\sigma = \gamma N^{-\frac{1}{4+d}}]

剪枝流程：从激活捕获到模型重建

第一阶段：数据准备

项目使用Hugging Face的ECE-ILAB/resilient-ai-unified数据集，通过scripts/download_dataset.py脚本下载并本地保存。这种设计允许在离线环境中进行剪枝操作，保护数据隐私。

第二阶段：激活捕获

scripts/prune.py脚本加载google/gemma-4-E4B-it模型（使用bfloat16精度以节省显存），通过PyTorch的hook机制捕获MLP层的down_proj激活。为了控制内存使用，脚本对token进行子采样处理。

第三阶段：冗余分析与聚类

互信息分数被转换为距离度量后，首先使用多维缩放（MDS）进行降维投影，然后应用KMeans聚类。在每个聚类中，仅保留距离质心最近的神经元，其余神经元被标记为冗余并从gate_proj、up_proj和down_proj权重矩阵中物理移除。

第四阶段：模型导出

剪枝后的模型连同更新后的Hugging Face配置（包括model.config.intermediate_size）一起保存，可直接用于后续推理任务。

技术实现细节

环境配置

项目采用uv作为包管理工具，安装简洁：

uv sync

PyTorch CUDA 12.4索引已在pyproject.toml中声明，无需额外配置。

运行流程

# 下载数据集
uv run python scripts/download_dataset.py

# 执行剪枝
uv run python scripts/prune.py

默认配置

MODEL_ID = "google/gemma-4-E4B-it"
DATASET_PATH = os.path.join(BASE_DIR, "data", "unified")
SAVE_PATH = os.path.join(BASE_DIR, "models", "gemma4-pruned-mi")

内存优化策略

使用device_map="auto"自动分配模型到可用GPU/CPU
bfloat16精度减少显存占用
激活数据的子采样处理
分块计算互信息矩阵

方法优势与适用场景

无监督特性

与传统剪枝方法不同，本项目无需标签数据即可识别冗余神经元。这使其特别适用于：

标注成本高昂的领域（如医疗、法律）
隐私敏感场景下的本地模型优化
快速原型验证和迭代

免重训练设计

剪枝后的模型可直接用于推理，无需经过耗时的微调过程。这一特性显著降低了模型压缩的工程门槛，使非专业用户也能受益于大模型压缩技术。

结构化剪枝的优势

与非结构化剪枝（稀疏矩阵）相比，结构化剪枝直接减少层宽度，带来：

更高的硬件友好性（稠密矩阵运算）
更直接的推理加速
更小的模型文件体积

局限性与改进方向

当前局限

剪枝比例需要人工设定，缺乏自适应机制
仅针对FFN层，未考虑注意力头的剪枝
互信息计算在神经元数量巨大时计算成本较高

潜在改进

引入重要性分数的动态阈值选择
扩展到注意力层和嵌入层的剪枝
探索更高效的核方法近似（如随机特征）
结合量化技术进一步压缩模型

项目结构与代码组织

Gemma4-Pruning/
├── main.py                    # 最小化入口点
├── scripts/
│   ├── download_dataset.py    # Hugging Face数据集下载
│   └── prune.py              # 核心剪枝逻辑
├── data/                     # 本地数据集存储
├── models/                   # 剪枝后模型输出
└── pyproject.toml            # 依赖与配置

代码结构清晰，职责分离明确，便于理解和扩展。

总结

Gemma4-Pruning项目展示了一种将信息论与深度学习相结合的大模型压缩新思路。通过互信息和Rényi熵量化神经元间的冗余性，该方法实现了无监督、免重训练的结构化剪枝，为Gemma 4等大语言模型的实际部署提供了可行的优化路径。

这一工作不仅具有工程实用价值，也为理解神经网络内部的信息流动提供了新的分析视角。随着大模型在边缘设备和资源受限场景中的部署需求日益增长，类似的压缩技术将发挥越来越重要的作用。