[论文解读] Adaptive wavelet distillation from neural networks through interpretations

Woo‐Seok Ha, Chandan Singh+3 more|arXiv (Cornell University)|Jul 19, 2021

Machine Learning in Materials Science参考文献 2被引用 20

一句话总结

本文提出自适应小波蒸馏（AWD），一种通过利用特征归因，将训练好的深度神经网络的预测知识蒸馏到可学习、可解释的小波变换中的方法。AWD生成了一个高度预测性、计算高效且具有科学可解释性的模型，在宇宙学参数推断和分子伙伴预测任务中优于最先进神经网络，同时识别出高密度峰值等与领域相关的特征，如收敛图中的高强度峰值和笼型蛋白动力学。

ABSTRACT

Recent deep-learning models have achieved impressive prediction performance, but often sacrifice interpretability and computational efficiency. Interpretability is crucial in many disciplines, such as science and medicine, where models must be carefully vetted or where interpretation is the goal itself. Moreover, interpretable models are concise and often yield computational efficiency. Here, we propose adaptive wavelet distillation (AWD), a method which aims to distill information from a trained neural network into a wavelet transform. Specifically, AWD penalizes feature attributions of a neural network in the wavelet domain to learn an effective multi-resolution wavelet transform. The resulting model is highly predictive, concise, computationally efficient, and has properties (such as a multi-scale structure) which make it easy to interpret. In close collaboration with domain experts, we showcase how AWD addresses challenges in two real-world settings: cosmological parameter inference and molecular-partner prediction. In both cases, AWD yields a scientifically interpretable and concise model which gives predictive performance better than state-of-the-art neural networks. Moreover, AWD identifies predictive features that are scientifically meaningful in the context of respective domains. All code and models are released in a full-fledged package available on Github (https://github.com/Yu-Group/adaptive-wavelets).

研究动机与目标

• 开发一种方法，以增强科学应用中深度神经网络的可解释性与计算效率。
• 通过基于小波的蒸馏，弥合高性能黑箱模型与内在可解释模型之间的差距。
• 通过与宇宙学和细胞生物学领域专家合作，将可解释性建立在真实科学问题的基础上。
• 构建一个兼具预测能力与科学意义的模型，识别与特定领域理论相关的特征。
• 证明基于归因蒸馏的小波模型在预测性能与压缩效率方面优于标准小波与深度神经网络。

提出的方法

• AWD利用预训练DNN的显著性特征归因，指导小波域中小波变换的学习。
• 通过在小波域中惩罚归因，优化小波变换，确保学习到的小波捕捉到与预测相关的关键特征。
• 通过施加重建误差与小波约束（如正交性、多分辨率结构），保持信号保真度与可解释性。
• 通过梯度下降在小波滤波器上进行优化，使用DNN的归因图作为监督信号。
• 所得小波变换为单层、自适应且可学习的变换，形成输入的稀疏多分辨率表示。
• 该方法支持小波激活图的可视化，可定位预测信息在输入中集中出现的位置。

实验结果

研究问题

• RQ1 我们能否将深度神经网络的预测知识蒸馏到一种既可解释又比原始模型更高效的基小波变换中？
• RQ2 如何利用教师DNN的归因信号，将基于小波的模型适配到数据与预测任务中？
• RQ3 AWD能否在真实科学问题中识别出具有科学意义的特征，如宇宙学中的峰值或笼型蛋白动力学？
• RQ4 蒸馏后的小波模型在预测性能与压缩效率方面是否优于标准小波或深度神经网络？
• RQ5 基于归因的优化如何提升小波特征与特定领域物理或生物机制之间的对齐性？

主要发现

• 在宇宙学参数推断任务中，AWD的均方根误差（RMSE）达到1.029 × 10⁻²，优于ResNet（1.156 × 10⁻²）与标准小波（1.569 × 10⁻²）。
• AWD小波模型的压缩率达0.610，优于DB5小波（0.620），实现了更高效的信号表示。
• 学习到的小波滤波器呈对称性，类似于宇宙学中用于峰值检测的匹配滤波器，表明其与已知物理原理一致。
• 小波激活图显示，AWD集中关注收敛图中的高强度峰值，证实了关于宇宙学信息含量的理论预期。
• 在分子伙伴预测任务中，AWD识别出笼型蛋白动力学为预测性特征，与已知的缓慢积累与衰减生物机制一致。
• 在两个科学领域中，AWD模型均提供了比标准小波与深度神经网络更简洁、可解释且更高效的替代方案。