PAHQ: Accelerating Automated Circuit Discovery through Mixed-Precision Inference Optimization
作者: Xinhai Wang, Shu Yang, Liangyu Wang, Lin Zhang, Huanyi Xie, Lijie Hu, Di Wang
分类: cs.LG, cs.AI
发布日期: 2025-10-27
🔗 代码/项目: GITHUB
💡 一句话要点
提出PAHQ,通过混合精度量化加速自动电路发现,提升大语言模型可解释性。
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 自动电路发现 机制可解释性 混合精度量化 激活修补 大语言模型
📋 核心要点
- 现有自动电路发现(ACDC)方法在应用于大型语言模型时,面临计算效率低和内存需求高的挑战。
- PAHQ的核心思想是利用激活修补与混合精度量化的对齐,对不同组件采用不同精度,提升效率。
- 实验表明,PAHQ加速ACDC可减少高达80%的运行时间和30%的内存消耗,同时保持分析的真实性。
📝 摘要(中文)
电路发现是机制可解释性的基石,旨在识别预训练语言模型中稀疏且与任务相关的子网络。自动电路发现(ACDC)已成为电路发现的关键方法,但其应用于大型语言模型时受到计算效率低下和内存需求过高的严重限制。虽然已经提出了一些加速方法,但它们主要依赖于ACDC的线性近似,这显著损害了解析的真实性。我们提出了一种加速自动电路发现的方法,即Per Attention Head Quantization (PAHQ),它通过优化每个单独的修补操作的效率来采取根本不同的方法。PAHQ利用了激活修补和混合精度量化(MPQ)之间的基本对齐:通过修补进行的可解释性分析本质上执行的是有针对性的消融研究。因此,我们可以专门为研究的组件保持高精度,同时安全地降低网络中其他位置的精度。与未加速的ACDC相比,PAHQ加速的ACDC将运行时间减少高达80%,内存消耗减少高达30%,同时保持了真实性。重要的是,我们的方法通过修改注意力计算机制,可以轻松地与现有的基于边缘的电路发现技术集成。这种无需训练的方法为加速机制可解释性方法提供了一种实用且新颖的途径。我们的代码可在https://github.com/626619403/PAHQ获得。
🔬 方法详解
问题定义:论文旨在解决自动电路发现(ACDC)应用于大型语言模型时计算效率低下和内存需求过高的问题。现有的加速方法通常采用线性近似,牺牲了解析的真实性,无法保证电路发现的准确性。
核心思路:论文的核心思路是利用激活修补和混合精度量化(MPQ)之间的内在联系。激活修补本质上是一种有针对性的消融研究,因此可以对需要高精度分析的关键组件保持高精度,而对其他组件降低精度,从而在保证分析质量的同时提升效率。
技术框架:PAHQ方法主要包含以下步骤:首先,确定需要进行电路发现的目标模型和任务。然后,在进行激活修补时,对不同的注意力头(Attention Head)采用不同的量化精度。具体来说,正在进行修补操作的注意力头保持较高的精度,而其他注意力头则可以降低精度。最后,通过实验验证PAHQ在加速ACDC的同时,是否能够保持电路发现的准确性。
关键创新:PAHQ的关键创新在于将混合精度量化与激活修补相结合,用于加速自动电路发现。与以往的线性近似方法不同,PAHQ通过选择性地降低精度,在保证分析真实性的前提下,显著提升了计算效率和降低了内存需求。这种方法无需训练,可以方便地集成到现有的电路发现流程中。
关键设计:PAHQ的关键设计在于如何确定不同注意力头的量化精度。论文中可能采用了一种启发式的方法,根据注意力头的重要性或者对任务的影响程度来决定其量化精度。此外,选择合适的量化方案(例如,量化比特数)也是一个重要的技术细节。具体的量化方案和精度选择策略需要在论文中进一步查找。
🖼️ 关键图片
📊 实验亮点
PAHQ方法在加速自动电路发现方面取得了显著的成果。实验结果表明,与未加速的ACDC相比,PAHQ可以将运行时间减少高达80%,内存消耗减少高达30%,同时保持了电路发现的准确性。这些结果表明PAHQ是一种高效且实用的加速方法,为大型语言模型的可解释性研究提供了新的途径。
🎯 应用场景
PAHQ加速自动电路发现的方法可广泛应用于大型语言模型的可解释性研究,帮助研究人员理解模型内部的运作机制,发现关键电路。这有助于提升模型的可信度,并为模型的改进和优化提供指导。此外,该方法还可应用于其他需要进行大规模消融实验的场景,例如神经网络架构搜索等。
📄 摘要(原文)
Circuit discovery, which involves identifying sparse and task-relevant subnetworks in pre-trained language models, is a cornerstone of mechanistic interpretability. Automated Circuit Discovery (ACDC) has emerged as a pivotal methodology in circuit discovery, but its application to large language models is severely limited by computational inefficiency and prohibitively high memory requirements. Although several accelerated approaches have been proposed, they primarily rely on linear approximations to ACDC, which significantly compromises analytical faithfulness. Our proposed method for accelerating automated circuit discovery, Per Attention Head Quantization (PAHQ), takes a fundamentally different approach by optimizing the efficiency of each individual patching operation. PAHQ leverages a fundamental alignment between activation patching and mixed-precision quantization (MPQ): interpretability analysis through patching essentially performs targeted ablation studies. Therefore, we can maintain high precision exclusively for investigated components while safely reducing precision elsewhere in the network. PAHQ-accelerated ACDC reduces runtime by up to 80\% and memory consumption by up to 30\% compared to unaccelerated ACDC while maintaining faithfulness. Importantly, our method readily integrates with existing edge-based circuit discovery techniques by modifying the attention computation mechanism. This training-free approach provides a practical and novel pathway for accelerating mechanistic interpretability methods. Our code is available at https://github.com/626619403/PAHQ.