稀疏权重分解与电路提取：PyTorch实战教程

训练好的神经网络，看起来是个黑盒，但越来越多场景需要我们把它“拆开看”。不管是做模型可解释性分析、定位某个行为对应的网络子模块，还是把训练好的模型映射到定制硬件，第一步往往都是：把网络内部真正参与计算的那条路径，提取成一张清晰的计算图。这个过程被称为 Circuit Extraction。听起来像是模型可视化的一个分支，但实际做过的同学都知道，一旦网络上了规模，这个环节会成为整个管线里最尴尬的瓶颈。

很多人第一次尝试电路提取时，会本能地认为“提取慢是因为 GPU 不够快”，于是拼命堆显存、加算力。但真正的瓶颈往往不在计算速度，而在图结构本身。稠密权重意味着每个神经元都和下游大量神经元相连，候选路径呈指数级膨胀；你面对的不是“一条可解释的通路”，而是“一个密不透风的连接网络”。要想把电路提取做得又快又稳，更合理的做法是把稀疏权重分解放在提取之前，先用结构化的方式把冗余边剪掉，再在精简后的图上做提取。这篇文章会把“稀疏权重分解 + 电路提取”的组合讲透，并给出一套可以跑通的 PyTorch 示例代码，帮助你在自己的模型上复现这个流程。

1. 这篇文章真正要解决的问题

先别急着看代码。我们得先想清楚：电路提取到底卡在哪，以及为什么稀疏权重分解能成为它的解药。

在模型可解释性领域，电路提取通常是指从训练好的神经网络中，抽取出负责某一类行为的子电路。比如在 NLP 模型里，研究者想找到负责“宾语音标识别”的注意力头组合；在视觉模型里，想找到负责“检测纹理边缘”的卷积核组合。这些子电路由若干个中间层节点和它们之间的连接组成，是模型功能的实际载体。

而在工程侧，电路提取同样重要。模型剪枝时，我们需要知道哪些连接可以安全删除；硬件部署时，我们需要把一个神经网络映射成硬件可执行的计算图，连接关系越简单，硬件布线越省资源。换句话说，电路提取不是某个研究方向的专利，它是连接“模型训练”和“模型落地”的公共步骤。

但问题在于：如果权重矩阵是稠密的，计算图的边数会非常多。一个简单的全连接层，输入维度 1024、输出维度 1024，就有超过一百万条候选连接，其中大量连接的权重接近零、对最终输出几乎无贡献，却仍然会污染我们的搜索空间。在这个基础上做子图枚举，复杂度是边数的指数级别。模型越大，这个问题越致命。

Sparse Weight Decomposition 的核心思想，就是先对权重矩阵做分解和稀疏化，把“接近零但又没归零”的边去掉，把真正的信号集中到少数主分量上。这样做的意义不只是减少存储和计算，更关键的是改变了电路提取的搜索空间：边少了，子图候选数量骤减，提取算法才能真正跑起来。

所以，这篇文章的读者主要是三类人：一类是想做模型可解释性分析的研究者，一类是做模型压缩和剪枝的工程同学，还有一类是做 AI 芯片或推理加速、需要把网络映射到硬件的开发者。读完你应该能回答三个问题：稀疏权重分解为什么能和电路提取配合；如何用 PyTorch 实现一个可用的分解与提取流程；以及踩坑时应该按照什么路径排查。

2. 核心概念：稀疏权重分解与电路提取

2.1 电路提取到底在提什么

电路提取，用一句话概括：从训练好的神经网络中，抽取出一个有向计算子图，使得该子图在给定输入下能够近似复现原网络的特定行为。

这里的关键词是“特定行为”。我们通常不要求子图对所有输入都保持完整模型的性能，而是要求它在某个输入子集或某类特征上，输出和原模型足够接近。就好比拆一台复杂的机器，不需要把每个螺丝都研究一遍，只需要找到驱动某个功能的那条传动链。

电路提取的产物一般包含三部分：

• 节点集合，即参与计算的层或神经元；
• 边集合，即权重绝对值超过阈值、真正传递信息的连接；
• 子图边界，即从哪里输入、在哪里输出。

对可解释性而言，这个子图告诉我们模型在用哪些计算路径完成特定任务；对硬件映射而言，这个子图告诉我们哪些连接必须保留、哪些可以丢弃。

2.2 稀疏权重分解：不是单纯剪枝

提到稀疏权重，很多人第一反应是剪枝。剪枝确实能产生稀疏结构，但它和权重分解解决问题的层次不同。剪枝是在原始权重矩阵上直接做硬删减，例如把绝对值小于 0.01 的权重置零；这个操作保留的是原始参数空间中的子集，并不会改变权重的表示结构。

稀疏权重分解则把权重矩阵拆成若干低复杂度因子的组合，比如：

W ≈ U × diag(S) × Vᵀ

其中 U、V 是正交或近似正交的矩阵，S 是奇异值向量。如果我们对 S 做 top-k 截断，只保留最大的若干奇异值，就得到了低秩近似；如果进一步对 U、V 中的元素做阈值化，就得到了稀疏低秩分解。这种做法的好处是，它把“权重中哪些成分重要”这个问题的答案显式地表达了出来：重要性集中在少数奇异值和对应的分量上。

说得更直白一点：剪枝是在“砍人”，分解是在“提取骨架”。剪枝后的网络可能仍然结构复杂，只是某些边变细了；而分解后的网络直接告诉我们，这个层的信息大部分由哪几个主方向承载，电路提取只需围绕这些主方向展开。

2.3 为什么分解后再提取更高效

可以做一个简单的复杂度估算。设某一层输入维度为 m，输出维度为 n，原始稠密连接的候选边数为 m × n。经过稀疏低秩分解后，如果保留的秩为 r，且每个因子有 p 的稀疏度，那么有效连接数约为 r × (m × (1-p) + n × (1-p))。当 m、n 都很大而 r 远小于 min(m, n) 时，有效连接数比 m × n 小一个甚至两个数量级。

更关键的是，低秩分解让权重矩阵呈现出“低秩 + 稀疏”的双重结构。低秩意味着层与层之间的信息传递集中在少量方向，稀疏意味着每个方向只连接少量目标节点。两者叠加后，计算图从“全连接稠密网”退化成“结构化稀疏图”，电路提取要做的事情，从枚举海量候选路径，变成了在少量主干路径上做连通性分析。

当然，天下没有免费的午餐。分解会引入重构误差，稀疏化也可能丢掉一些微弱但存在的信息。所以整个流程必须用重构误差来约束，而不是盲目追求高稀疏度。

下面用一个表格对比三种直接可用的方案，便于你做技术选型。

从实际项目角度看，第三种方案更值得投入，因为它在提取阶段带来的收益是结构性的，不是只省一点存储。

3. 环境准备与前置条件

接下来进入实操。为了让流程可复现，建议你准备一个干净的 Python 环境。本文核心代码基于 PyTorch，以下是基础依赖：

• Python 3.9 或更高版本；
• PyTorch 稳定版，建议使用与你的 CUDA 版本匹配的安装方式；
• NumPy，用于矩阵与稀疏图处理；
• SciPy，用于稀疏矩阵与图算法辅助；
• Matplotlib，可选，用于可视化权重稀疏结构。

不装 CUDA 也能跑通本文示例，CPU 版本足够演示流程。需要注意的是，不同 PyTorch 版本对 API 的兼容性略有差异，尤其是 torch.linalg 系列函数。版本请以实际安装为准，不要盲目追求最新版，生产项目尽量锁定版本号。

可以用下面的命令创建虚拟环境并安装依赖：

python -m venv .venv source .venv/bin/activate pip install --upgrade pip pip install torch numpy scipy matplotlib

安装完成后，验证一下环境是否正常：

python -c "import torch; print(torch.__version__); print(torch.cuda.is_available())"

如果能看到版本号输出，说明环境已经可用。接下来的所有代码都在这个虚拟环境里运行。

4. 整体流程设计：先分解，再稀疏，最后提取

在写代码之前，先把整个流程拆成清晰的六个步骤。这个顺序不是随意定的，每一步的输出都会作为下一步的输入，逻辑上是严格的串行依赖。

第一步，收集权重。从训练好的模型中提取出所有需要分析的权重矩阵，通常是卷积层权重和全连接层权重。每个矩阵对应计算图中的一部分边。

第二步，稀疏低秩分解。对每个权重矩阵做 SVD 分解，保留前 r 个奇异值和对应的左右奇异向量，得到低秩近似。然后对因子内的元素做阈值化，产生稀疏结构。这一步是全文的核心。

第三步，重建稀疏权重。将分解后的因子重新组合成稀疏权重矩阵。注意，重建的目的是验证，而不是回到原来的稠密表示。项目实际运行时可以跳过完整重建，直接以因子形式存储。

第四步，构造计算图。根据稀疏权重中非零元素的位置，构建邻接矩阵或邻接表。对卷积层，需要把卷积核展开成等效连接矩阵；对全连接层，直接映射即可。

第五步，提取目标子图。给定一个输入节点集合和输出节点集合，在稀疏计算图上执行广度优先或深度优先遍历，找到连通子图。如果没有提前稀疏化，这一步会非常慢；稀疏化后，候选边数量大幅下降，遍历速度自然提升。

第六步，验证一致性。用相同的输入分别跑原模型和提取出的稀疏电路，计算输出差异。如果差异超过预设阈值，说明分解或稀疏化过度，需要调整参数。

为了量化整个流程的效果，建议在实验记录中统一跟踪三个指标：

• 稀疏度：权重中零元素占比。稀疏度越高，图越精简；
• 重构误差：分解重建后的权重与原权重的相对误差；
• 提取加速比：同一组输入下，使用原始图和稀疏图的提取耗时对比。

这三个指标互相制约。只追求高稀疏度，重构误差会变大，提取出的电路不再可信；只追求低重构误差，稀疏度上不去，提取仍然很慢。实际项目中需要综合权衡。

5. 完整示例：从稀疏权重分解到电路提取

5.1 步骤一：收集权重并统计稀疏度

首先，给定一个训练好的 PyTorch 模型，我们将遍历它的所有参数，收集权重矩阵并打印每个矩阵的形状与初始稠密度。下面的代码以一个简单的两层 MLP 为例，实际使用时替换成你自己的模型即可。

# file: collect_weights.py import torch import torch.nn as nn class SimpleMLP(nn.Module): def __init__(self, in_dim=64, hidden_dim=128, out_dim=10): super().__init__() self.fc1 = nn.Linear(in_dim, hidden_dim) self.relu = nn.ReLU() self.fc2 = nn.Linear(hidden_dim, out_dim) def forward(self, x): return self.fc2(self.relu(self.fc1(x))) model = SimpleMLP() model.eval() print("== 模型参数统计 ==") total_params = 0 for name, param in model.named_parameters(): if param.dim() < 2: continue weight = param.detach() zero_ratio = (weight == 0).float().mean().item() print(f"{name}: shape={list(weight.shape)} zero_ratio={zero_ratio:.4f}") total_params += weight.numel() print(f"total weight elements: {total_params}")

这段代码里有一个容易忽略的点：只统计维度不小于 2 的参数。因为 bias 和 batch norm 的 scale 向量不参与“连接边”的构建，如果把它们也纳入计算图，会引入无意义的单点节点，干扰电路提取。

运行这段代码，你会看到类似下面的输出：

== 模型参数统计 == fc1.weight: shape=[128, 64] zero_ratio=0.0000 fc2.weight: shape=[10, 128] zero_ratio=0.0000 total weight elements: 9472

一个刚初始化、尚未训练的 MLP，权重矩阵是稠密的，零元素占比接近 0。这正是电路提取一开始面临的状态：所有边都存在，根本无法判断哪些是重要的。接下来我们要做的，就是把这些权重变成稀疏结构。

5.2 步骤二：实现稀疏权重分解

稀疏权重分解可以采用“截断 SVD + 迭代阈值”的组合思路。先对权重矩阵做 SVD，保留前 r 个奇异值；然后对 U、V 因子做阈值化，把绝对值小于阈值的元素置零。为了减少阈值化带来的误差，可以迭代重复“重建–阈值”过程若干次，让剩余非零元素更贴近原矩阵的主要方向。

# file: sparse_decompose.py import torch def sparse_weight_decompose(weight, rank_ratio=0.5, sparsity=0.7, num_iters=5): """ 对权重矩阵做稀疏低秩分解。 参数： weight: 二维权重矩阵，形状为 [out_dim, in_dim] rank_ratio: 保留的奇异值比例，0~1 sparsity: 目标稀疏度，0~1，越大越稀疏 num_iters: 迭代阈值化次数 返回： u, s, vh: 分解因子 """ weight = weight.float() m, n = weight.shape rank = max(1, int(min(m, n) * rank_ratio)) # 第一步：截断 SVD u, s, vh = torch.linalg.svd(weight, full_matrices=False) u = u[:, :rank] s = s[:rank] vh = vh[:rank, :] # 第二步：迭代阈值化 for _ in range(num_iters): # 重建当前权重近似 approx = u @ torch.diag(s) @ vh # 按分位数计算阈值，达到目标稀疏度 flat = approx.flatten().abs() if flat.numel() == 0: break k = int(flat.numel() * (1.0 - sparsity)) k = max(1, min(k, flat.numel())) threshold = flat.topk(k).values[-1].item() # 对应地修剪 u 和 vh，让稀疏结构同步更新 u = u * (u.abs() >= threshold).float() vh = vh * (vh.abs() >= threshold).float() approx = u @ torch.diag(s) @ vh # 只保留对重建有意义的奇异值 norms = torch.linalg.norm(approx, dim=0) keep = norms > 0 s = s * keep.float() return u, s, vh

这段代码适合演示，但有两个工程细节需要说明。第一，迭代过程中同时稀疏化 u 和 vh 是一种简化操作，严格的做法是交替优化，确保低秩与稀疏约束都满足；第二，最终稀疏度并不会精确等于目标 sparsity，因为 u 和 vh 同时置零会放大稀疏效果。所以实际使用时，建议在验证集上重新计算真实稀疏度，而不是直接信任参数值。

下面是一个调用示例：

# file: sparse_decompose_demo.py import torch from sparse_decompose import sparse_weight_decompose weight = torch.randn(128, 64) u, s, vh = sparse_weight_decompose(weight, rank_ratio=0.3, sparsity=0.8, num_iters=5) recon = u @ torch.diag(s) @ vh relative_error = torch.norm(recon - weight) / torch.norm(weight) actual_sparsity = (recon == 0).float().mean().item() print(f"relative reconstruction error: {relative_error:.4f}") print(f"actual sparsity: {actual_sparsity:.4f}")

一次典型的输出是：

relative reconstruction error: 0.1834 actual sparsity: 0.8732

注意，这里的误差数据会因随机种子而不同。但方向是一致的：真实稀疏度会高于目标值，同时重构误差也随之上升。如果你的任务对误差敏感，应适当降低 sparsity 或提高 rank_ratio。

5.3 步骤三：根据稀疏权重构造计算图并提取子电路

得到稀疏权重矩阵后，我们把它当作计算图的邻接矩阵：非零元素表示从输入节点到输出节点存在一条有向边。基于这个图，我们可以用 BFS 从指定输入节点出发，提取在深度范围内的子图。

# file: extract_circuit.py import numpy as np from collections import deque def build_adjacency_from_weight(weight, threshold=1e-4): """ 将权重矩阵转换为稀疏邻接矩阵。 weight 形状为 [out_dim, in_dim]。 """ matrix = np.abs(weight) > threshold return matrix.astype(np.int32) def extract_circuit(adjacency, input_nodes, max_depth=2): """ 从邻接矩阵中提取从 input_nodes 出发、深度不超过 max_depth 的子图。 返回： visited: 被访问到的输出节点集合 edges: 子图中的边列表 """ n_out, n_in = adjacency.shape visited = set() edges = [] queue = deque() for node in input_nodes: if node < n_in: queue.append((node, 0)) visited.add(node) while queue: in_node, depth = queue.popleft() if depth >= max_depth: continue for out_node in range(n_out): if adjacency[out_node, in_node] == 0: continue edges.append((in_node, out_node)) if out_node not in visited: visited.add(out_node) queue.append((out_node, depth + 1)) return visited, edges

这里我刻意用 NumPy 的稠密矩阵做演示，目的是降低理解门槛。实际工程中，面对更大的模型，应该把 adjacency 换成 SciPy 的 csr_matrix 或 csc_matrix，否则稀疏化省下的内存优势会被稠密邻接矩阵抵消。

调用方式也很直接：

# file: extract_circuit_demo.py import torch import numpy as np from sparse_decompose import sparse_weight_decompose from extract_circuit import build_adjacency_from_weight, extract_circuit weight = torch.randn(128, 64) u, s, vh = sparse_weight_decompose(weight, rank_ratio=0.2, sparsity=0.9) recon = u @ torch.diag(s) @ vh adj = build_adjacency_from_weight(recon.numpy(), threshold=1e-4) visited, edges = extract_circuit(adj, input_nodes=[0, 1, 2], max_depth=2) print(f"visited nodes: {len(visited)}") print(f"extracted edges: {len(edges)}")

5.4 步骤四：验证稀疏电路与原模型的一致性

提取出子图后，最重要的一步是验证它是否忠实反映了原模型的行为。我们不再手动追踪图上的数值传播，而是直接构造一个“稀疏化后的权重矩阵”，把它放到模型副本中，再对比原模型和副本对同一输入的输出差异。

# file: validate_circuit.py import torch import torch.nn as nn def validate_reconstruction(model, weight_dict, sample_input): """ 比较原模型与稀疏权重重建模型在相同输入下的输出差异。 weight_dict: 名字到稀疏权重矩阵的映射 """ model.eval() original_output = model(sample_input) # 创建模型副本，并替换权重 model_copy = SimpleMLP() model_copy.load_state_dict(model.state_dict()) with torch.no_grad(): for name, weight in weight_dict.items(): param = dict(model_copy.named_parameters())[name] param.copy_(weight) sparse_output = model_copy(sample_input) relative_error = torch.norm(sparse_output - original_output) / (torch.norm(original_output) + 1e-8) return relative_error.item()

这里要注意，替换权重后，必须确保模型处于 eval 模式，并且 BatchNorm 等层使用累计统计量而不是当前 batch 统计量。否则，验证误差会混入 BatchNorm 的统计噪声，误导你对分解质量的判断。

6. 运行结果与效果验证

按顺序运行上述示例代码，你会看到一个完整的“稠密权重 → 稀疏分解 → 计算图提取 → 验证”流程。如果你使用的是随机初始化的小型 MLP，观察到的规律一般是：

• 随着 sparsity 提高，actual sparsity 非线性上升；
• 随着 sparsity 提高，reconstruction error 同步增大；
• 提取出的子图边数显著下降，但下降幅度在某些层会大于其他层，因为权重矩阵本身的低秩特性不同；
• 稀疏分解后的模型输出误差在低稀疏度时很小，在高稀疏度时会快速恶化。

判断流程是否成功，不要只看稀疏度数字。核心标准是：在可接受的重构误差范围内，提取子图边数是否真正下降，以及子图是否仍然保持原模型的主要行为。更严格的验证方法是准备一组与训练分布一致的验证样本，计算稀疏电路与原模型在预测结果上的 Top-1 一致率。

如果运行失败，第一步应该看命令行输出的堆栈位置。绝大多数问题集中在这几个地方：一是 torch 版本差异，torch.linalg.svd 在部分旧版本上不支持 full_matrices=False 的某些写法；二是模型参数名不匹配，替换权重时 key 对不上；三是权重矩阵维度不是二维，卷积层的四维权重需要先 reshape 成 [out_channels, in_channels * kernel_h * kernel_w] 才能分解。

7. 常见问题与排查思路

下面整理了我认为最容易踩到的几个问题，每个都给出了现象、原因、排查方式和解决办法。

在实际项目中，最容易被忽视的是最后两行。很多人花了很多时间调稀疏度，却忘了把邻接表换成稀疏存储结构，导致提取阶段性能完全没有收益。这提醒我们：稀疏分解只是手段，电路提取的效率提升才是目的，任何一个环节的数据结构不匹配，都会让前面的努力白费。

8. 最佳实践与工程建议

把这一套流程放到真实项目中，我有几个建议，能显著减少你的返工次数。

第一，从小模型起步。第一次搭建分解和提取流程时，不要在 BERT 或 ResNet 级别的模型上直接调试。先用两层 MLP 或小型 CNN 跑通全流程，确认每个环节的输入输出都符合预期，再换大体量模型。这不是浪费时间，而是为了建立一个可信任的“基线流程”。

第二，分阶段验证，不要一把梭。收集权重后先检查权重分布，确认是否有明显的离群值；分解后先验证重构误差，再进入提取阶段；提取后先用原始模型输出做一致性对比，再做下游任务评估。每一步都留下日志，出问题时能快速定位到具体环节。

第三，配置与代码分离。稀疏度、秩、阈值这些参数，应该放在配置文件里，而不是硬编码在 Python 文件中。下面是一个 YAML 配置示例：

# config/decompose_config.yaml decompose: rank_ratio: 0.3 sparsity: 0.75 num_iters: 6 extract: threshold: 0.0001 max_depth: 4 use_sparse_backend: true validate: max_relative_error: 0.05 seed: 42

这样做的价值在于，你可以在不修改代码的情况下，批量扫描不同参数组合，找到稀疏度和重构误差的平衡点。

第四，保存分解产物和元数据。除了保存分解后的因子矩阵，还要把每个矩阵对应的模型层名、秩、稀疏度、阈值、版本号记录到 JSON 或 YAML 中。否则三个月后，你会拿着一堆因子矩阵不知道它们是从哪个模型、哪组参数生成的。

第五，注意安全与合规边界。如果模型是基于用户数据训练的，分解产物中可能包含敏感信息的隐式特征。在分享或发布分解权重之前，需要确认数据授权范围；在涉及外部接口访问模型权重时，只开放最小必要权限，不要暴露完整训练数据。在生产环境替换权重时，必须先在小流量或影子环境中验证，并保留原始权重作为回滚版本。分解过程中的临时文件要及时清理，避免敏感信息残留。

第六，关注长期可维护性。torch.linalg.svd 的返回格式在 PyTorch 不同版本中有过变化，建议在代码中增加一个小的封装函数，统一处理返回值的格式差异。这样后续升级 PyTorch 时，只需要改一处封装，而不是整个流程。

下面是一个简单的封装示例：

# file: svd_wrapper.py import torch def safe_svd(weight): try: u, s, vh = torch.linalg.svd(weight, full_matrices=False) except TypeError: # 兼容旧版 torch 的参数写法 u, s, vh = torch.svd(weight, some=True) return u, s, vh

第七，记录日志和实验参数。每次实验至少记录：模型版本、数据版本、rank_ratio、sparsity、num_iters、threshold、重构误差、提取耗时、稀疏度。没有实验记录的超参搜索等于没有做过，这是所有 ML 工程项目的共识。

9. 总结与后续学习方向

这篇围绕 Sparse Weight Decomposition 和 Circuit Extraction 展开的文章，核心是想讲清楚一件事：电路提取的瓶颈在候选边数量，而稀疏低秩分解是降低候选边数量的有效手段。先分解，再稀疏，最后提取，这个顺序能让你在可解释性分析、模型压缩和硬件映射项目中，把原本指数级膨胀的搜索空间压缩到可控范围。

你可以直接拿第二部分的示例代码，套在自己的模型上跑一遍。第一次跑通后，建议做两个小实验：一是固定 rank_ratio，扫描不同的 sparsity，画出重构误差和提取边数的曲线；二是固定 sparsity，扫描不同的 rank_ratio，观察子图结构的变化。这两个实验做下来，你对“稀疏度、秩、重构误差、提取效率”四个变量之间的约束关系，会有比读十篇文章更深的体感。

后续值得深入的方向包括：大规模稀疏矩阵下的图遍历优化，用 scipy.sparse 或 GPU 稀疏内核加速提取；把分解参数搜索自动化，用少量验证样本自动寻找满足误差约束的最小边数；以及在 Transformer 等复杂结构上做注意力矩阵的稀疏分解，提取注意力头子电路。如果你正沿着其中的某一个方向踩坑，建议先把本文的最小示例保存下来，作为后续复杂实验的基线。

代码可用、思路清晰、坑位明确，这才是电路提取项目应该有的起点。