用最优传输方法快速定位神经网络中的因果变量,提升可解释性研究效率。
PLOT: Progressive Localization via Optimal Transport in Neural Causal Abstraction

- 基于最优传输建立抽象变量与神经单元的全局对应关系
- 单层神经元级定位即可达成高精度,大模型分步细化定位
- 显著提速传统方法,适合大规模神经网络因果分析
因果抽象为机制可解释性提供了严谨框架,通过反事实干预分析将高层因果模型与神经网络的底层计算对齐。现有方法如分布式对齐搜索(DAS)虽能学习表达性强的子空间干预,但相关神经位置事先未知,需在候选位置上进行计算量巨大的搜索才能找到干预点。本文提出PLOT(基于最优传输的渐进式定位),通过抽象与神经干预输出效应几何结构之间的最优传输耦合,实现因果变量的全局软定位,可转化为干预操作符。在简单场景下,单个神经元层级的耦合即足够;在大型模型中,PLOT逐步推进,从令牌、时间步或层等粗粒度单元,到坐标组或主成分跨度等细粒度支持集,并可选地引导DAS。在多层级复杂实验中,仅使用传输的PLOT干预速度极快且精度相当;而由PLOT引导的DAS仅需全量DAS一小部分运行时间即可达到其精度水平,为大规模因果抽象研究提供高效定位引擎。
原文摘要 · Abstract (English)
Causal abstraction offers a principled framework for mechanistic interpretability, aligning a high-level causal model with the low-level computation realized by a neural network through counterfactual intervention analysis. Existing methods such as distributed alignment search (DAS) learn expressive subspace interventions, but the relevant neural site is unknown a priori, so finding a handle requires a computationally burdensome search over candidate sites. We introduce PLOT (Progressive Localization via Optimal Transport), a transport-based framework that localizes causal variables from the output effect geometry of abstract and neural interventions. PLOT fits an optimal transport coupling between abstract variables and candidate neural sites, yielding a global soft correspondence that can be calibrated into intervention handles. In simple settings, a single coupling over individual neurons suffices. In larger models, PLOT is applied progressively, moving from coarse sites such as tokens, timesteps, or layers to finer supports such as coordinate groups or PCA spans, and optionally guiding DAS based on the localized signal. Across experiments of increasing complexity, transport-only PLOT handles are exceedingly fast and competitive on accuracy, while PLOT-guided DAS reaches DAS-level accuracy at a fraction of full DAS runtime, providing an efficient localization engine for causal abstraction research at scale.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。