safety-research/selective-gradient-masking

# 选择性梯度掩码 (SGTM) 论文 **"Beyond Data Filtering: Knowledge Localization for Capability Removal in LLMs"** 的配套代码。 本仓库提供了 **选择性梯度掩码 (SGTM)** 的实现，这是一种通过参数级别的梯度操作来定向移除语言模型特定能力的方法。代码包含了论文实验中使用的训练脚本、模型实现和分析 notebook。 ## 概述 SGTM 通过以下方式实现语言模型特定能力的选择性遗忘： - 将参数专用于目标能力 - 在训练过程中掩码梯度，仅更新相关参数 - 在训练后移除（消融）专用参数 ## 安装说明 ``` uv pip install -e . ``` ## 运行测试 测试是使用 pytest 实现的： ``` pytest ./src -sv ``` ## 代码结构 ``` ├── src/sgtm/ │ ├── model/ # SGTM model implementations and masking strategies │ │ └── tests/ # Unit tests for model components │ ├── train/ # Training scripts for different experiments │ └── data/ # Data processing and tokenization scripts ├── configs/ │ ├── tinystories/ # Model hyperparameters for TinyStories experiments (8M, 18M, 34M, 64M) │ └── wiki/ # Model hyperparameters for Wikipedia experiments (34M, 64M, 125M, 254M) ├── scripts/ │ ├── tinystories/ # Bash scripts for TinyStories experiments │ ├── wiki/ # Bash scripts for Wikipedia experiments │ └── data/ # Data preparation scripts └── notebooks/ # Analysis and visualization notebooks ``` ## 掩码策略 本仓库实现了三种互补的选择性梯度掩码方法： ### 1. **参数掩码**（论文中的 SGTM） 论文中使用的主要方法。参数被拆分为“遗忘”（forget）和“保留”（retain）维度。在保留数据上训练时，反向传播后会将与遗忘维度相关的梯度置零，从而防止更新这些参数。 **实现文件**：`parameter_masking.py` ### 2. **梯度路由** ([Cloud et al., 2024](https://arxiv.org/abs/2410.04332)) 基于 Cloud 等人的方法，该方法在激活级别而非参数级别进行操作。在前向传播过程中使用 `.detach()` 来阻断通过特定激活维度的梯度流，从而防止梯度到达相应的参数。 **实现文件**：`gradient_routing.py` ### 3. **激活掩码** 在保留数据上进行前向传播时，将与遗忘维度对应的激活置零，从而有效地仅通过保留参数来路由信息。 **实现文件**：`activation_masking.py` ## 模型实现 ### 使用 SGTM 初始化模型 SGTM 扩展了 HuggingFace Transformers 中的 GPT-Neo 架构。以下是创建具有选择性梯度掩码功能的模型的方法： ``` from transformers import GPTNeoConfig from sgtm.model import GPTNeoForCausalLMSGTM # 使用 SGTM 参数配置模型 config = GPTNeoConfig( vocab_size=50257, hidden_size=512, num_layers=12, num_heads=32, max_position_embeddings=2048, # SGTM-specific parameters retain_mlp_dim=1984, # MLP dimensions for retain data (out of 2048 total) retain_attn_heads=31, # Attention heads for retain data (out of 32 total) masking_strategy="parameter_masking", # Strategy to use split_masked_weights=True, # Use SplitLinearOut for keeping retain and forget weights in separate parameters (no effect on traininig) sgtm_mask_embeddings=False, # Whether to mask embedding gradients during training ) # 初始化模型 model = GPTNeoForCausalLMSGTM(config) # 在训练期间，指定你要训练的 data split outputs = model( input_ids=batch["input_ids"], labels=batch["labels"], sgtm_mode="forget" # or "retain" or "default" ) # 在 backward pass 之后，根据 data split 调整梯度 loss.backward() model.adjust_gradients(sgtm_mode="forget") optimizer.step() ``` ### 核心类 - **`GPTNeoForCausalLMSGTM`**：支持 SGTM 的主模型类 - **`GPTNeoModelSGTM`**：带有选择性掩码的基础 transformer - **`GPTNeoBlockSGTM`**：带有掩码的独立 transformer 块 - **配置参数**： - `retain_mlp_dim`：专用于保留数据的 MLP 维度数量 - `retain_attn_heads`：专用于保留数据的注意力头数量 - `masked_layers`：应用掩码的层索引列表 - `masking_strategy`：指定使用的掩码策略 - `split_masked_weights`：是否物理拆分参数 ## 实验 ### TinyStories 实验 TinyStories 实验旨在从基于双语 TinyStories 训练的模型中移除某种语言的知识。 #### 数据集 双语（英语-西班牙语）TinyStories 数据集可在 HuggingFace 上获取： - **数据集**：[`ffuuugor/tinystories-spanish`](https://huggingface.co/datasets/ffuuugor/tinystories_spanish)（派生自 [roneneldan/TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories)） - **翻译**：该数据集最初为英文。西班牙语翻译是使用 `src/sgtm/data/translate_to_spanish.py` 创建的 #### 数据准备 在运行实验之前，请对数据进行分词和拆分： ``` bash scripts/data/tinystories_tokenize.sh ``` 这将： 1. 从 HuggingFace 下载该数据集 2. 分离英语和西班牙语版本 3. 使用 GPT-2 tokenizer 对两者进行分词 4. 将其保存到 `data/datasets/tinystories_split/en`（保留）和 `data/datasets/tinystories_split/es`（遗忘） #### 运行实验 可以使用 `scripts/tinystories/` 中的脚本启动论文中的关键实验： ``` # 在 TinyStories 上使用 SGTM 训练模型，并改变未发现 forget data 的比例 bash scripts/tinystories/mislabeling.sh # 在 TinyStories 上使用 SGTM 训练模型，同时跟踪每个 checkpoint 的 calibrated loss 以进行 trade-off 分析。 bash scripts/tinystories/tradeoff.sh ``` #### 分析 Notebook - **`notebooks/tinystories/retain_forget_tradeoff.ipynb`**：生成二维遗忘/保留权衡图，比较 SGTM、梯度路由和数据过滤基线 - **`notebooks/tinystories/undiscovered_rate.ipynb`**：分析未发现的遗忘数据（错误标记）对模型性能的影响 - **`notebooks/tinystories/grad_norms.ipynb`**：计算并可视化不同数据拆分下的遗忘和保留参数的梯度范数 ### Wikipedia 实验 #### 数据准备 [ORES 主题数据](https://www.mediawiki.org/wiki/ORES/Articletopic)可在[此处](https://analytics.wikimedia.org/published/datasets/topics/)下载。 然后我们通过以下方式仅筛选英文文章： ``` cat topicsForAllWikipediaPages2020-08-24.csv | grep "enwiki," > enwiki_topics2020.csv ``` 准备用于实验的 Wikipedia 数据集： ``` bash scripts/data/prepare_wikipedia.sh ``` 此脚本将： 1. 从 HuggingFace 下载 Wikipedia 数据集 2. 将 Wikipedia 文章与 ORES 主题数据进行匹配 3. 使用 GPT-2 tokenizer 对数据进行分词 4. 将文章拆分为上下文大小（1024 个 token）的块 5. 将数据拆分为三个独立的数据集： - **遗忘** (Forget)：Stem.Biology 文章（需要移除的目标能力） - **邻近遗忘** (Forget-adjacent)：化学、医学和环境文章（相关主题） - **通用知识** (General knowledge)：所有其他主题 6. 每个类别预留 5000 篇文章用于测试集 #### 运行实验 可以通过以下命令启动论文中的关键实验： ``` bash scripts/wiki/run.sh ``` 该脚本会启动 4 次运行： - SGTM - 弱过滤 - 严格过滤 - 无过滤 额外的基线方法： ``` # RMU (Representation Misdirection for Unlearning) bash scripts/wiki/rmu.sh # unlearning 之后进行 Finetuning bash scripts/wiki/finetune.sh ``` `rmu.sh` 脚本实现了 RMU 基线方法，该方法通过引导模型表示来执行机器遗忘。`finetune.sh` 脚本可用于在机器遗忘后对模型进行微调，以评估鲁棒性。 #### 分析 Notebook - **`notebooks/wiki/retain_forget_tradeoff.ipynb`**：生成二维权衡图，比较遗忘损失（生物学）与保留损失（文化/地理/历史）以及邻近遗忘损失（医学/化学/环境） - **`notebooks/wiki/finetuning.ipynb`**：分析微调期间的模型性能，比较不同的机器遗忘方法（RMU、过滤器、SGTM）在持续训练下的维持能力 ## 引用 ``` (Citation will be added upon paper publication) ```

标签：Apex, DLL 劫持, NoSQL, 人工智能, 凭据扫描, 大语言模型, 安全规则引擎, 机器学习, 梯度掩码, 模型训练, 用户模式Hook绕过, 知识遗忘, 逆向工具