稀疏权重分解:神经网络电路提取实战指南

发布时间:2026/8/28 15:47:31
稀疏权重分解:神经网络电路提取实战指南 做模型解释的人应该都遇到过这种局面想从一个训练好的神经网络里抽出几条决定性计算路径结果打开权重一看全是稠密矩阵。每个输入神经元都和每个隐藏神经元连着每个隐藏单元又和全连接输出连着画出来的“电路”跟蜘蛛网没有区别既无法人工检查也没办法拿去验证。稀疏权重分解做的事情就是先把这种稠密连接压成一个真正可以阅读的稀疏子图再从这个子图出发做电路提取。这个概念在模型可解释性里越来越常见。所谓电路提取并不是把神经网络还原成硬件电路而是找到网络中“哪些神经元、哪些连接、沿着什么顺序完成了一个具体行为”的子结构。稀疏权重分解是在这个流程最前面的降噪步骤先去掉不重要的连接保留少数主干边然后基于主干图做因果验证。这样提取出来的结果才具备可解释性也更容易和别人交流。下面按我的实操习惯拆开讲。我会先用一个小型布尔逻辑模型做验证对象再给出完整流程、参数判断标准和常见坑。如果你的目标是一个真实 transformer流程主体可以复用差异主要在规模控制和逐层处理上。1. 为什么电路提取要先做稀疏权重分解1.1 稠密权重无法直接画出电路训练好的神经网络权重矩阵几乎都是稠密的。哪怕某个连接权重很小小到对输出只有千分之一的影响它在权重矩阵里依然占一个位置。直接拿这个矩阵画图最终结果是全连接图。全连接图带来的问题很直接无法定位关键路径所有输入到所有输出都有“连接”。可视化没有区分度边的粗细和颜色掩盖了结构。后续做激活修补或因果追踪时候选边太多实验组合爆炸。所以我一般不会直接把原始权重拿去做图分析。第一件事是把每个权重矩阵先拆解再稀疏化用极少数保留边代替完整矩阵。1.2 稀疏权重分解解决哪一步不解决哪一步稀疏权重分解解决的是“候选结构生成”这一步。它把网络内部可能性极高的连接数量降下来让后续可以针对少数子图做精细验证。但它不解决“这个子图是否真的负责了某个行为”。权重大小只能代表相关性不能直接等同于因果性。一个连接权重很大不代表它一定驱动了某个预测。所以完整流程必须是对权重做低秩分解和稀疏化。从稀疏矩阵转成候选电路图。用扰动、激活修补或因果追踪验证子图有效性。很多人在这里跳步把“稀疏后的权重图”直接当成“模型内部电路”这是不严谨的。稀疏图只是候选集经过验证之后才有资格叫提取出来的电路。2. 环境准备和实验对象2.1 依赖和硬件条件下面的流程用 PyTorch 就可以复现。建议环境Python 3.9 或更高。PyTorch 2.x。NumPy。可选的 networkx用于画图和导出边列表。硬件方面如果只是验证小型 MLPCPU 就够了。如果换到小型 transformer建议准备至少 8 GB 显存并把序列长度控制在 128 到 512 之间。下面这个示例根本不涉及大规模激活保存内存压力很小。pip install torch numpy networkx这里有一个容易忽略的点PyTorch 的低秩分解接口在不同版本之间有小差异。旧版本常用torch.svd新版本更推荐torch.linalg.svd。测试前先确认自己的 torch 版本避免接口不一致。2.2 用一个小型布尔逻辑模型当验证对象我建议第一次做实验不要直接上大语言模型。流程没跑通之前拿大模型只会增加调试成本。先用一个结构简单、答案可人工核对的小模型。我构造了一个三输入布尔函数label (x1 AND x2) OR (x3 AND NOT x1)输入三位分别是 0 或 1输出也是 0 或 1。这个函数本身有清晰逻辑x1 同时控制两条可能路径x2 和 x3 分别作为后续条件。如果电路提取成功我们应该能观察到这两条并行路径且看到 x1 对两条路径的抑制作用。模型用一个两层 MLPimport torch import torch.nn as nn class ToyCircuit(nn.Module): def __init__(self): super().__init__() self.fc1 nn.Linear(3, 16) self.act nn.ReLU() self.fc2 nn.Linear(16, 1) def forward(self, x): return self.fc2(self.act(self.fc1(x)))训练时加入一点 L1 正则化让权重尽量稀疏。因为原始逻辑本身是稀疏的一个能拟合函数的网络通常也不需要太多冗余连接。model ToyCircuit() optimizer torch.optim.Adam(model.parameters(), lr1e-3) loss_fn nn.BCEWithLogitsLoss() x torch.randint(0, 2, (4096, 3)).float() y ((x[:, 0] * x[:, 1]) (x[:, 2] * (1 - x[:, 0]))).clamp(max1).float() for epoch in range(300): optimizer.zero_grad() logits model(x) loss loss_fn(logits, y.unsqueeze(1)) l1 sum(p.abs().sum() for p in model.parameters()) total_loss loss 1e-4 * l1 total_loss.backward() optimizer.step()训练完成后网络预测准确率应该接近 100%。这时再进入分解步骤。2.3 工作流程拆成四段整个流程我习惯拆成四段避免一上来就写一坨代码低秩分解对每个线性层权重做 SVD。稀疏掩码根据阈值或比例生成稀疏矩阵。图构建把非零连接转成有向图。因果验证对候选图做扰动和一致性检查。后面每段都单独验证不要一把梭。3. 稀疏权重分解实操流程3.1 读取权重并做低秩分解以fc1为例权重形状是16 x 3。这个矩阵很小直接做 SVD 很快。W model.fc1.weight.detach() U, S, Vh torch.linalg.svd(W, full_matricesFalse)这里U的每一列是一组输出方向Vh的每一行是一组输入方向S是奇异值。奇异值越大说明该方向在权重中越重要。低秩分解的意思就是只保留前 k 个方向丢掉后面大部分小方向。k 3 W_low U[:, :k] torch.diag(S[:k]) Vh[:k, :]k 的选择没有固定值要结合矩阵维度看。对16 x 3的矩阵秩最多是 3所以 k 到 3 基本就够。对真实 transformer一个4096 x 1024的权重矩阵k 可以取 64 或 128具体看重建误差。重建误差可以用这个公式衡量relative_error torch.norm(W - W_low) / torch.norm(W)我自己一般把相对误差控制在 5% 到 20% 之间。太低说明还保留了太多冗余太高说明丢失了有效结构。3.2 构造稀疏掩码的三种方式低秩分解只能降低秩不能直接把矩阵变稀疏。接下来还要做掩码。第一种是绝对阈值掩码。设定一个阈值tau小于阈值的权重直接置零。这个方式最直观但问题是对权重尺度敏感。不同层的权重分布差异很大同一套阈值可能在 A 层去掉 10% 的边在 B 层去掉 99% 的边。第二种是比例掩码。把所有权重按绝对值排序只保留最大的 top p 比例。这个方式更容易控制最终稀疏度。第三种是标准差掩码。以权重的均值加若干倍标准差作为阈值。这种方法对分布自适应更强。我在小型模型上常用的是一套组合先做低秩截断再做标准差掩码。sigma W_low.std() mask (W_low.abs() 1.2 * sigma).float() W_sparse W_low * mask这一段的含义是只有那些在低秩重建后仍然显著偏离零的连接才进入候选电路。注意不要一上来就追求极端稀疏比如只保留 1% 的边。图漂亮不等于结果正确。我的经验是先把稀疏度控制在 10% 到 30%验证能通过后再进一步压缩。3.3 把稀疏矩阵转成电路子图拿到W_sparse之后可以把它当成邻接矩阵来看。行是目标神经元列是源神经元非零元素就是一条有向边。edges torch.nonzero(W_sparse, as_tupleFalse) for row, col in edges.tolist(): weight W_sparse[row, col].item() print(f{col} - {row} weight{weight:.4f})如果愿意还可以转成 networkx 的有向图方便后续分析和画图import networkx as nx G nx.DiGraph() for row, col in edges.tolist(): G.add_edge(col, row, weightW_sparse[row, col].item())这里要特别提醒一个神经网络的真正“电路”是多层组合不是只看单个矩阵。fc1的非零边决定哪些输入进入哪些隐藏神经元fc2的非零边决定哪些隐藏神经元影响输出。两层合起来才能构成“输入到输出”的完整路径。所以流程应该是对fc1和fc2分别做低秩分解和掩码。合并两个矩阵的非零边。找出同时出现在两端路径中的隐藏神经元。只保留这些神经元作为候选电路节点。这一步很关键。否则你只是画了一个稀疏权重矩阵不是一条完整电路。4. 用因果验证确认稀疏图有效4.1 重建输出和一致性检查稀疏图只是候选结构。判断它是不是有效电路最终要看它能不能复现原始模型的行为。最简单的一致性检查是这样把原始模型权重替换成稀疏化之后的权重然后输入所有可能的输入组合看输出和原始模型相差多少。def forward_with_sparse(model, x): h torch.relu(x model.fc1.weight.T model.fc1.bias) return (h * mask_fc2) model.fc2.weight.T model.fc2.bias注意要同时掩码fc2否则来自“非候选神经元”的信息还是会穿过网络。对三输入布尔函数全部组合只有 8 个。可以直接逐条对比原始模型预测为 1。稀疏模型预测也为 1。原始模型预测为 0。稀疏模型预测也为 0。如果所有样本一致说明候选电路基本覆盖了关键行为。如果某些样本翻转不一致说明掩码砍掉了必要边需要先回去调参数不要急着下结论。4.2 扰动验证看翻转是否对齐一致性检查只是基础。更严格一点要做扰动验证。把某个隐藏神经元置零观察原始模型输出变化然后同样把稀疏图对应节点置零观察稀疏模型输出变化。如果两个模型表现出一致的翻转方向说明这条路径确实承担了相似功能。def zero_neuron(model, layer_index, neuron_idx, x): with torch.no_grad(): if layer_index 0: h model.act(model.fc1(x)) h[:, neuron_idx] 0 return model.fc2(h) else: h model.fc2(model.act(model.fc1(x))) h[:, neuron_idx] 0 return h对比时主要看三件事扰动哪些神经元会导致输出发生变化。改变的方向是一致的还是相反的。改变幅度是否相近。如果稀疏图提取正确两个模型的“敏感神经元集合”应该高度重合。4.3 先单样本再全样本强烈建议先固定一个输入样本做扰动验证不要一开始就在全部样本上跑。比如固定输入为x11, x21, x31这时函数真值来自x1 AND x2路径。再看哪些神经元被置零后输出会翻转。单样本验证通过后再扩大到全部 8 个输入组合。这样能快速排除“掩码选择导致某条路径意外消失”的问题。5. 参数、结果判断和效率对比5.1 关键参数速查表下面是我在实验中最常调的一组参数放在一起方便对照检查。参数作用初值建议调整方向低秩保留数 k控制低秩重建精度矩阵维度较小的 1/3重建误差高时增大相对重建误差判断 k 是否够用5% 到 20%误差太大时增大 k掩码阈值倍数控制稀疏度1.0 到 1.5 倍标准差断开时降低倍数稀疏边比例衡量最终候选边规模10% 到 30%链路不完整时提高比例扰动神经元因果验证的范围从 3 到 5 个开始确认稳定后再全量这些参数在真实 transformer 上要按层单独统计不能一套参数全模型共用。因为 attention 层的权重和 MLP 层的权重分布差异很大。5.2 结果判断标准怎么判断提取出来的电路好坏我一般用四条标准路径闭合。从某个输入神经元出发能找到一条通往输出的连续路径。行为对齐。稀疏模型和原始模型在测试样本上输出一致或高度一致。扰动响应一致。对关键节点做置零扰动两个模型表现方向一致。子图规模可控。保留边数比全连接图低一个数量级以上。如果只有前两条满足说明稀疏图还只是“结构近似”还没到“可解释电路”的程度。要四条都满足才比较稳。5.3 效率和资源占用怎么比很多人会把“稀疏权重分解”理解成“模型推理加速”。这是误解。分解和掩码本身要额外耗时不会让原模型跑得更快。它的效率价值在后续分析阶段。假设全连接图中候选边有 1 万条激活扰动实验要做 1 万次组合测试。如果稀疏化后只保留 1000 条实验次数直接少一个数量级。对这种“分析效率”做对比才合理。测量时建议记录分解耗时。掩码耗时。候选边数量。扰动验证每组平均耗时。全部验证集完成耗时。不推荐只用墙钟时间对比因为环境差异大。更好的指标是“验证同样的行为需要多少次扰动”这个数字和硬件无关。6. 常见问题和排查链路6.1 稀疏图分裂成碎片如果保留边太少或者阈值太高稀疏图经常会变成几个互不相连的小组件。这时从输入到输出没有一条完整路径。排查顺序先看低秩重建误差如果误差超过 30%说明 k 太小。再看阈值倍数降低阈值倍数。检查是否同时掩码了所有层单看某一层容易误判。如果只是学习验证不要追求 1% 以下的极端稀疏度把边比例放到 15% 到 30% 更实用。6.2 稀疏图输出与原模型偏差大常见原因不是稀疏化本身而是掩码把关键边切掉了。排查顺序先只做稀疏化不做低秩看输出偏差多大。再只做低秩不做稀疏看偏差多大。如果两者单独都没问题组合才出问题说明阈值倍数和 k 的配合不合适。输出偏差很大时优先降低低秩 k而不是调整掩码阈值。因为低秩重建误差会放大后续掩码的影响。6.3 换成大模型后计算变慢怎么办真实 transformer 的权重矩阵大得多。直接对整个权重矩阵做完整 SVD 会很慢尤其是 attention 里面的投影矩阵。这时有三个降级方案使用torch.svd_lowrank做截断 SVD只求前 k 个奇异向量。按层处理不要一次加载所有权重。对 attention 和 MLP 分开构建子图再合并公共节点。我通常会先把 MLP 的权重分解做完再处理 attention 的 query、key、value 和 output 投影。否则合并后的图太复杂很难做因果验证。6.4 低配置机器的降级方案如果你的机器没有 GPU甚至内存都比较小第一步可以先不用 transformer。就用小型 MLP 或者单层 attention 的 toy model 跑通整套流程把阈值的规律摸清楚。下面这套方案适合低配环境使用不超过 4 层的模型。输入维度控制在 8 以内。隐藏层神经元控制在 64 以内。不保存全量激活只保存稀疏权重矩阵和边列表。可视化用 networkx 导出边列表不要直接在图上渲染几万条边。低配环境也能验证核心思路只是不能用它作为生产级可解释性分析的最终结论。7. 最后再留几个我自己会优先看的点稀疏权重分解用于电路提取真正决定成败的不是 SVD 本身而是后续验证环节。权重稀疏化只能帮你缩小搜索范围不能替代因果实验。所以每跑完一次提取我都会问自己三个问题这个子图里有一条从输入到输出的连续路径吗把某条边切掉后模型行为是否按照预期方向变化换一个随机种子重新训练提取出来的子图是否稳定第三个问题最容易被忽略。神经网络训练有随机性不同训练轮次拟合同一个函数可能形成不同内部结构。如果提取出的电路结构只在某一次训练里成立那它不能算模型的普遍规律。稳妥做法是训练多个种子提取多份候选图找它们的公共骨架。我个人更建议先把单模型、单任务的稀疏图和因果验证跑稳再往大型模型上迁移。流程跑通后稀疏权重分解反而会成为你分析模型内部行为时最常用的前置步骤。它不酷但非常实用。