Lesson 1: MLLM-SAE

本课要训练的是一个读取冻结模型内部表示的 SAE。图片、问题和已有答案先经过 MLLM；SAE 学习重构其中一个位置的 hidden state，并把它表示成少量 feature 的组合。下面的交互图是可操作的教学模拟，feature 编号、名称与激活值均不是实测结果。

本课以 Qwen/Qwen3.5-9B 和 google/gemma-4-E4B-it 为例，目标是为图文指令数据建立可解释的去重与覆盖指纹。先回答最关键的问题：一个模型、一个 hook、一个共同训练的 SAE，可以让图片位置和文本位置共用一个 65,536 维字典；两个模型不能共用同一个 SAE。字典也没有预先划分“图片区”和“文本区”。

模型与实现核对日期：2026-11-06。E4B 的模型链接是 google/gemma-4-E4B-it，不是 E2B。Qwen 家族的 GitHub 地址虽然已经指向 Qwen3.8，本课固定使用 Qwen3.5-9B checkpoint。本文的层位、数据配比与验收门槛是待验证的实验起点，没有论文证明它们在这两个模型上最优。

01 · 跟随一张图片，走进模型

图中用一张包含两个红苹果、一个杯子和 SALE 字牌的教学插画，配合问题与答案。切换模型、hook 或 token 来源，观察张量形状与 SAE 输入如何改变；点击“播放一次”沿着前向计算逐步查看。图片不会经过文本 tokenizer；进入语言主干的是视觉编码后的向量，占据序列中的图片位置。

主实验选 C：语言主干完整 block 的输出。这个位置保留了残差流中的综合信息，适合同时研究视觉内容、提问方式与答案信息。选择它是因为研究目标是整条图文指令的数据选择，而不是单独解释某个 attention 或 MLP 分支。

三种 hook，各自回答不同问题

这里的路径相对于 Hugging Face 的完整 ForConditionalGeneration 对象，层号均从 0 开始。config 给出尺寸与层数；模块路径要结合实现源码确认，启动时再用 named_modules() 和一次前向校验。以下形状来自配置与源码推导，本文没有运行这两个模型的权重。Qwen 配置 · Gemma 配置 · Qwen 实现 · Gemma 实现

“完整 block 输出”不等于 mlp.down_proj 输出。 Qwen 的 block 输出已经加回 residual；Gemma 还包含 MLP 后归一化、每层输入分支（PLE）与 layer scalar。若 hook 在 down_proj 或 self_attn.o_proj，得到的只是分支贡献，不能沿用 C 的解释与统计。

两个主干的差异，为什么会影响 hook

示例图片为正方形：Qwen 若处理后为 512×512，则有 1,024 个视觉 patch，经 merger 得到 256 个视觉 token。Gemma 按当前 processor 的 max_soft_tokens=280 处理正方形时，resize 到 768×768：2,304 个真实 patch，pad 到 2,520；pooling 得到 280 个槽，去掉 24 个 padding 槽后，实际有 256 个视觉 token。280 是这里的预算，不保证每张图都有 280 个有效 token。 长宽比变化时数量也会变化，代码必须读实际 mask，不能硬编码图中的 256。Gemma processor 源码

02 · 模型与 SAE，各自吃什么、输出什么

d 是一个 token 向量的坐标数，不是概念数；第 17 个 hidden 维度不天然代表“红色”。SAE 希望把原来叠加在这些坐标里的信息，分解为更容易检验的字典方向。但单个 feature 仍可能多义、碎片化，或者只反映位置与格式。

图片位置也在语言主干里，但不因此变成文本

对一条 图片 → 用户问题 → 已有答案 的样本，给每个位置保留来源标签：

以上因果解释适用于本课的图在前、问题在后排序及这两个 checkpoint 的配置。Gemma E4B 的 use_bidirectional_attention 为 null，不要把其他 Gemma 变体的图像双向 mask 套过来。视觉编码器内部可全局读取图片，不代表 LLM 图片位置可以读取未来的问题。

本课的单图样本统一把首轮 user 内容规范为 [image, question_text]，再使用官方 chat template；保留正文和真实轮次，并记录这一排序规则。部分来源把问题文本放在图片之前，直接沿用会让前置 Q token 看不到图片，不能套用上表。多图或顺序具有特殊含义的样本留待独立方案处理。

不把所有位置不加区分地喂给 SAE。 丢弃 padding、system、角色/轮次/图片边界等模板 token；保留真实 I/Q/Ans。图片较多并不应自动获得更高权重。先按来源与图片抽样，再在各池中限额均匀抽 token。真实多轮保留对话历史；同一图片附带的独立 VQA 问答应拆成独立样本，不能拼成虚假的对话。

Ans 是否参与取决于目标：本课筛选的是完整 SFT 样本，所以参与；若只研究“输入图片 + 问题”的表示，就去掉 Ans，并用相同条件重训/重评。用于验证最终任务的参考答案不能混入 SAE 训练集。teacher forcing 只是读已有训练答案，不会自动让错误答案变正确。

03 · 用什么数据：先覆盖目标分布

SAE 训练集是产生 hidden 的输入语料；SAE 的重构目标则由冻结模型提供，不需要人工 feature 标签。以 LLaVA 风格的单图指令数据为目标，采用“广泛图文概念 + 属性/关系问答 + 少量 OCR + 指令格式”的组合。这里先做英文图文；面向中文或医学、图表等新领域时，要替换相应份额并重新验证，不能宣称这份配比普遍最优。

每个模型、每个主 SAE 的正式上限为 200M 次有效激活暴露，先用包含在其中的 20M 做 pilot。一次暴露是一条 [d] 向量进入一次 SAE 更新；它不是原始文本 token 数，也不是完整 MLLM forward 的序列 token 数。重复抽样计入暴露预算；同时记录独立图片数、独立文本 token 数与 forward 总量。Qwen 与 Gemma 各训一个主 SAE，总上限 400M，A/B 可选实验另计。

表中镜像是仓库访问入口，不保证当前位置的下载速度或所有文件直连；部分请求会重定向到 Hugging Face。固定数据 revision、split、过滤规则与许可证记录。H4 的混合数据不是原始 LLaVA-Instruct-150K。原始 LLaVA-Instruct-150K 的标注还需要配套 COCO 图片，不能只下载对话 JSON 就开始抽视觉激活。

来源配比与 token 池配比是两层抽样。 在每个来源中，先以 I 50% / 文本 50% 为起点；问答来源的文本份额在 Q 与 Ans 中等分，caption 来源的文本份额给真实 caption/Ans，排除人为添加的固定“Describe this image”请求。短文本不能靠反复重复几个 token 填满预算，应扩大图片/样本覆盖并报告实际比例。无需额外加纯文本，图文样本已有大量文本位置；纯文本是后续泛化对照，不能给 A/B 产生视觉激活。

入库前完成：坏图/空文本检查 → 统一 RGB 与官方 template → 图片 ID、精确 hash 和近重复分组 → 排除最终评测图片及其近重复 → 按图片组切分 90/5/5 train/validation/test，seed=42。COCO、VQAv2、LLaVA 等来源可能复用图片，不能按行随机切分；某些第三方 train split 也包含 COCO val 图片。额外报告与候选筛选池的重叠；以未见图片/来源的评估检验泛化。

04 · 从原始数据到训练更新

每个 token 独立成为一个 SAE 训练样本，混合后打散。源图片和对话不会丢失：它们的 ID、token 位置与来源标签随缓存保存，后续可以回看某个 feature 激活最强的原图区域与文本。

冻结 MLLM → forward → hook → 采样 → 固定归一化 → SAE → 重构 loss → 仅更新 SAE。 SAE 并不串接在正常生成路径中，不要求 MLLM 学会使用它；冻结模型每次提供原始 hidden target。不要在 no_grad() 内运行需要更新的 SAE，只有 MLLM 的激活采集放在其中。

训练配方：先跑得对，再扩大预算

定义 D ∈ R^(d×K) 为 decoder 字典，encoder 为 W ∈ R^(K×d)。每个 D[:,j] 是第 j 个 feature 的重构方向；z_j 是当前 token 使用它的强度。

为使 AuxK 的实现没有歧义：e=stop_gradient(x−x_hat)；只在 dead 集合内从 ReLU(u) 选取最多 k_aux 个系数，记为 z_dead；辅助重构为 D z_dead，不加 decoder bias。

无 dead feature 时令辅助项为 0；初始 bias 为 0。TopK 已控制稀疏度，本配方不再叠加一个未经调参的 L1 惩罚。这个归一化辅助 loss 是本课选定的实现约定；不要把它与其他 SAE 仓库的不同 reduction、残差归一化或 batch TopK 默认值混用。TopK、dead-feature 处理及重构/行为多维评估的依据可参见 Scaling and evaluating sparse autoencoders。

65K 不意味着训练代价小：两块稠密权重约为 2×d×K 参数。Qwen SAE 约 5.37 亿参数，Gemma SAE 约 3.36 亿；FP32 参数、梯度及 Adam 两份矩状态合计约 8.6 GB / 5.4 GB（十进制），尚未计主模型、activation、临时张量与框架开销。先测显存；不足时减 microbatch、累积梯度，或将 32K 作为明确记录的独立配置。65K 在两个模型上的扩张倍数分别是 16× 和 25.6×，不能据此视为容量完全匹配的科学对照。

可运行的 hook 烟雾检查

下面代码只验证“模型确实经过该模块、形状正确、image mask 与序列对齐”，不是完整训练器。需要可加载模型的 GPU 环境、支持这两个模型的 Transformers 版本，以及 Gemma 已获准的模型访问。先对安装版本运行，再锁定版本与模型 revision；本文没有执行权重加载。

pythonCopyimport torch
from PIL import Image
from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration, Gemma4ForConditionalGeneration

which = "qwen"  # 改成 "gemma"
specs = {
    "qwen": ("Qwen/Qwen3.5-9B", Qwen3_5ForConditionalGeneration, 23, 4096),
    "gemma": ("google/gemma-4-E4B-it", Gemma4ForConditionalGeneration, 29, 2560),
}
model_id, cls, layer, d = specs[which]
processor = AutoProcessor.from_pretrained(model_id)
model = cls.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto")
model.eval().requires_grad_(False)
path = f"model.language_model.layers.{layer}"
modules = dict(model.named_modules())
assert path in modules, path
print(path, type(modules[path]).__name__)

image = Image.open("example.jpg").convert("RGB")
messages = [
    {"role": "user", "content": [
        {"type": "image", "image": image},
        {"type": "text", "text": "What is in this image?"}]},
    {"role": "assistant", "content": [
        {"type": "text", "text": "Replace this with the dataset's actual answer."}]},
]
batch = processor.apply_chat_template(
    messages, tokenize=True, add_generation_prompt=False,
    return_dict=True, return_tensors="pt",
).to(model.device)
cache = {}
def capture(_module, _args, output):
    h = output[0] if isinstance(output, tuple) else output
    assert h.ndim == 3 and h.shape[-1] == d
    cache["h"] = h.detach()

handle = modules[path].register_forward_hook(capture)
try:
    cache.clear()  # 每次 forward 都清空，避免误用上一样本的激活
    with torch.no_grad():
        _ = model.model(**batch, use_cache=False)
    h = cache.pop("h")
    ids = batch["input_ids"].to(h.device)
    valid = batch["attention_mask"].to(h.device).bool()
    image_mask = ids.eq(model.config.image_token_id) & valid
    assert h.shape[:2] == ids.shape
    assert image_mask.any(), "图文样本应存在真实 image slots"
    print("all:", tuple(h.shape), "image:", tuple(h[image_mask].shape))
    # Q / Ans 用模板渲染时记录的角色正文跨度对齐，不能用“非 image”代替。
    # 抽样后再离开 no_grad，用 FP32 SAE 计算 loss/backward/optimizer.step。
finally:
    handle.remove()

生产数据处理需要补齐 Q/Ans mask：记录 chat template 中每个消息正文的 token span，与 processor 展开 image slots 后的序列对齐；支持 assistant mask 的模板可利用该能力，但要检查 mask 是否全零、是否覆盖真实答案。用单图、纯文本、独立 QA、真实多轮和 padding 样本逐个检查。不要通过 decoded 文本里搜索答案字符串定位，重复词与特殊 token 会造成错位。对齐完成后，assert 三个池互斥、均不含 padding，且所有入池位置都有来源记录。

第一轮限定单图、文本上限约 512 token、总序列上限 2048；超过上限的样本显式跳过并记录比例，不截断半个图像 span。Qwen 固定 processor 的像素预算（如 65,536–1,048,576）；Gemma 固定 max_soft_tokens=280 起步；两者保留长宽比，缓存与正式筛选使用相同配置。OCR 受分辨率影响，后续提高预算必须重验表示分布。

05 · 一个 65K 字典，里面到底装什么

字典不是预先命名的概念表，也不是按比例分给图片/文本的内存。同一个 C hook 上的 I、Q、Ans 向量，共享同一套 encoder、decoder 和 65,536 个 feature ID。 抽样时 50% 图像位置并不意味着训练后恰好 50% feature 属于图片。

点击 feature 查看它在三类位置上的示意激活。改变 k 只改变每个 token 的稀疏预算，不改变字典槽位数。图中把 65,536 个槽缩略为一个 256×256 网格；点亮的是合成示例，数值不构成实际模型的 feature 发现。

训练时不分配语义类别；训练后才统计。 对每个 feature 记录 I/Q/Ans 的激活概率、强度、source 分布与 top examples，按池样本数归一化；允许多个标签、未解释与噪声类别。需要跨模态语义一致性时，增加图文对应、图像替换、问题改写和 feature 干预实验。VL-SAE 使用专门的显式对齐与模态 decoder 来推动一致概念；普通混合 TopK SAE 不自动获得同样性质。VL-SAE

一轮建议抽查 100 个活跃 feature，每个回看至少 20 个不同图片组/样本的 top examples，并包含随机激活与负例。对主 hook，固定同图与前缀、只改后置问题时，I 应在数值容差内保持不变，而 Q/Ans 可以改变；固定问题换图、遮挡相关区域，则检验文本位置的 feature 是否确实依赖视觉。

06 · 怎样拿这些特征筛数据

先对所有有效 token 逐个编码成 z，再聚合。先平均 hidden 再经过 SAE 会改变非线性稀疏编码的结果。对每条样本分别保存 v_I, v_Q, v_Ans：各池 mean(z) 后 L2 归一化用于相似度；另外保存超过 feature 阈值的激活集合，用于覆盖。

TopK 的每 token 非零数接近固定，因此“激活总次数多”不是多样性分数。 相似度阈值与覆盖阈值在 held-out 校准集上选择，并检查稀有任务与回答长度的偏差。给每个 feature 保存数据源证据；新覆盖只在活跃、稳定且排除明显模板/噪声的集合中计分。

07 · 什么结果才算可以使用

不要只看总 MSE。总体 loss 很好可能只是图片位置占多数，把问题与答案的重构问题掩盖了。

FVU 用 Σ||h−h_hat||² / Σ||h−mean(h)||²；按 I/Q/Ans 分别算，低方差池单独报告绝对误差，避免不稳定比值。所有数值门槛都是诊断起点，最终能否用于选择数据，要由误删率、任务保留与下游 SFT 决定。

训练结束至少保存：SAE 权重与 bias、μ/s、K/k、模型与 processor revision、Transformers 版本、完整 hook 路径、数据 split/hash/采样配置、实际预算、按池验证指标、feature 使用统计与 top-example 索引。没有这些信息，一个权重文件很难被正确复用。

08 · 哪些来自研究，哪些是本课的选择

读完本课后，应能把四件事分开：输入图片/文本是什么，冻结模型产生的 hidden 是什么，SAE 学到的字典方向是什么，数据选择使用的样本指纹是什么。 下一步是先用 20M pilot 检查这些边界，再决定是否扩大预算。