基于 Transformer RoBERTa的情感分类任务实践总结之五——剪枝

注意力头重要性剪枝（Head Importance Pruning）

这种剪枝方法旨在从像 RoBERTa 这样的预训练 Transformer 模型中移除最不重要的注意力头，同时不显著降低其性能。其核心思想是量化每个注意力头对模型整体损失的贡献，然后移除贡献最小的那些头。

算法总结

1. 初始化注意力头掩码（Head Mask）：

   • 创建一个 head_mask 张量，并设置 requires_grad=True。这个掩码的维度与模型中层数和每层注意力头的数量相同。
   • 初始时，掩码中的所有值都设为 1，这意味着所有注意力头都是活跃的。

2. 计算注意力头重要性：

   • 将模型设置为评估模式 (model.eval())。
   • 对于给定的一批数据，执行模型的正向传播，并应用 head_mask。head_mask 作为注意力头输出的乘数，让模型“看到”哪些头正在被考虑，是权重调节器。
   • 计算损失。
   • 关键步骤是调用 loss.backward()。这一步会计算损失对 head_mask 的梯度。这些梯度的绝对值可以作为每个头重要性的代理——更大的梯度表明改变该头的输出（通过掩码）对损失有更大的影响，从而暗示其重要性更高。
   • 梯度被分离（detach）并按层内注意力头进行归一化，以获得相对重要性分数。

3. 识别要剪枝的注意力头：

   • 对于每一层，识别出重要性分数最低（由计算出的梯度决定）的注意力头。
   • 将这些最不重要的头收集到一个字典 (heads_to_prune) 中，该字典将层索引映射到要剪枝的头索引列表。

4. 剪枝注意力头：

   • 调用 model.prune_heads() 方法，并传入 heads_to_prune 字典。该方法会内部修改模型的架构，以移除指定的注意力头。这通常涉及调整后续层（例如输出投影层）的维度，以适应减少的注意力头数量。

5. 微调：

   • 剪枝后，通常建议在您的数据集上对剪枝后的模型进行微调。这使得剩余的头和模型其他参数能够适应剪枝后的架构，并可能恢复任何性能损失。在您的代码中，模型在剪枝后立即被传递给 Trainer 进行训练。

为什么它有效？

这种方法之所以有效，是因为它直接量化了每个注意力头对模型主要目标（最小化损失）的贡献。通过移除贡献最小的头，实际上是在降低模型的冗余和计算复杂性，同时旨在保留其最关键的信息通路。

代码

import torch
from transformers import (
    RobertaForSequenceClassification,
    RobertaTokenizer,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset
import os
from huggingface_hub import snapshot_download
from torch.utils.data import DataLoader
import numpy as np
from sklearn.metrics import accuracy_score, f1_score

os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

# 下载模型到本地
model_path = snapshot_download(
    repo_id="roberta-base",
    local_dir="./local_roberta_model",
    local_dir_use_symlinks=False,
)
print(f"模型已下载到: {
     
     model_path