Experiment

本项目包含用于大语言模型（LLM）的 Sparse Autoencoder (SAE) 实验代码，主要用于分析模型在拒绝回答（refusal）任务中的特征激活模式。

项目结构

主要代码文件

• sea_Gemma.py: Gemma 模型的 SAE 探针训练和评估脚本
• sea_llama.py: Llama 模型的 SAE 探针训练和评估脚本
• sea_Qwen.py: Qwen 模型的 SAE 探针训练和评估脚本
• sea_llama_stage2.py: Llama 模型的第二阶段实验脚本
• get_data.py: 数据获取和处理脚本
• get_data_llama.py: Llama 模型专用数据获取脚本
• make_data.py: 数据制作脚本
• get_feature.py: 特征提取和筛选脚本
• keshihua.py: 可视化脚本
• sae_safe_loader.py: 安全的 SAE 加载器
• sae.py: SAE 相关工具函数
• setup_hf.py: HuggingFace 环境配置脚本
• demo.py: 演示脚本

配置文件

• requirements.txt: Python 依赖包列表
• install_deps.sh: 依赖安装脚本
• model_sae_map.json: 模型和 SAE 映射配置文件

文档

• 实验记录.md: 实验参数和配置记录

安装

1. 安装依赖

# 使用安装脚本（推荐）
bash install_deps.sh

# 或手动安装
pip install -r requirements.txt

2. 配置 HuggingFace

运行 setup_hf.py 配置 HuggingFace 缓存路径（可选）。

使用方法

训练 SAE 探针

# Gemma 模型
python sea_Gemma.py \
    --data_path /path/to/data \
    --model google/gemma-2-9b-it \
    --release gemma-scope-9b-pt-res \
    --output_dir /path/to/output

# Llama 模型
python sea_llama.py \
    --data_path /path/to/data \
    --model /path/to/model \
    --release llama_scope_lxr_8x \
    --output_dir /path/to/output

# Qwen 模型
python sea_Qwen.py \
    --data_path /path/to/data \
    --model Qwen/Qwen2.5-7B-Instruct \
    --release qwen-scope-7b-instruct \
    --output_dir /path/to/output

提取特征

python get_feature.py \
    --model_name gemma-2-9b-it \
    --map_json model_sae_map.json

数据处理

# 获取数据
python get_data.py

# 制作数据
python make_data.py

依赖

主要依赖包括：

• PyTorch >= 2.0.0
• Transformers >= 4.35.0
• sae-lens
• scikit-learn
• numpy (>= 1.24.0, < 2.0.0)

完整依赖列表请参见 requirements.txt。

注意事项

• 本项目使用 HuggingFace 缓存，默认缓存路径为 /root/autodl-tmp/Projects
• 数据和实验产出目录（*_data/, *_sae/）不会被提交到 Git 仓库
• 确保有足够的 GPU 内存来运行实验

许可证

本项目仅供研究使用。