[论文解读] Hardware Acceleration of Explainable Machine Learning using Tensor Processing Units

Zhixin Pan, Prabhat Mishra|arXiv (Cornell University)|Mar 22, 2021

Advanced Neural Network Applications参考文献 10被引用 4

一句话总结

本论文提出首个基于TPU的可解释机器学习硬件加速框架，利用矩阵卷积与傅里叶变换，将复杂的解释任务转化为高效的矩阵运算。相较于基于GPU的方法，该框架在分类任务中实现最高38.6倍的加速，在解释时间上实现12.8倍的加速，从而在多种机器学习工作负载中实现结果的实时解释。

ABSTRACT

Machine learning (ML) is successful in achieving human-level performance in various fields. However, it lacks the ability to explain an outcome due to its black-box nature. While existing explainable ML is promising, almost all of these methods focus on formatting interpretability as an optimization problem. Such a mapping leads to numerous iterations of time-consuming complex computations, which limits their applicability in real-time applications. In this paper, we propose a novel framework for accelerating explainable ML using Tensor Processing Units (TPUs). The proposed framework exploits the synergy between matrix convolution and Fourier transform, and takes full advantage of TPU's natural ability in accelerating matrix computations. Specifically, this paper makes three important contributions. (1) To the best of our knowledge, our proposed work is the first attempt in enabling hardware acceleration of explainable ML using TPUs. (2) Our proposed approach is applicable across a wide variety of ML algorithms, and effective utilization of TPU-based acceleration can lead to real-time outcome interpretation. (3) Extensive experimental results demonstrate that our proposed approach can provide an order-of-magnitude speedup in both classification time (25x on average) and interpretation time (13x on average) compared to state-of-the-art techniques.

研究动机与目标

• 解决可解释机器学习（XML）中解释过程缓慢的关键瓶颈，从而支持实时部署。
• 克服现有XML方法依赖迭代、耗时的优化问题所带来的低效性。
• 通过利用TPU在加速矩阵计算方面的原生优势，实现实时结果解释。
• 开发一个可泛化的框架，适用于多种机器学习算法与领域，如图像分类与恶意软件检测。
• 通过硬件感知的算法设计，在分类与解释两个阶段均实现显著加速。

提出的方法

• 将可解释机器学习中的模型蒸馏过程映射为线性代数问题，转化为一系列矩阵运算。
• 利用矩阵卷积与傅里叶变换之间的协同效应，简化并加速解释计算。
• 利用TPU的脉动阵列架构——特别是其256×256矩阵乘法单元（MXU）——实现超高速的8位整数矩阵乘法。
• 应用数据分解技术，将大矩阵拆分为更小的子矩阵，实现跨多个MXU的高效分布，降低带宽压力。
• 将模型蒸馏与TPU加速的前向传播相结合，生成输入特征的贡献因子，实现结果解释。
• 设计框架以在TPU硬件上将分类与解释统一为单一、优化的流水线。

实验结果

研究问题

• RQ1 使用TPU进行硬件加速是否能显著减少可解释机器学习中结果解释所需的时间？
• RQ2 是否可以将XML中的模型蒸馏过程重构为可在TPU上高效执行的矩阵计算？
• RQ3 所提出的基于TPU的框架在矩阵规模与模型复杂度增加时的可扩展性如何？
• RQ4 所提出的方法在分类与解释速度方面，相较于基于GPU的最先进技术，能多大程度上实现超越？
• RQ5 该框架是否能在保持高性能的同时，提供可解释且局部化的解释（例如，识别关键特征如恶意软件攻击模式）？

主要发现

• 所提出的基于TPU的框架相较于最先进的基于GPU的方法，平均实现25倍的分类时间加速。
• 解释时间平均加速13倍，且在CIFAR-100基准测试中达到最高39.5倍的加速。
• 对于1024×1024的矩阵，所提方法比基线快逾30倍，展现出强大的可扩展性。
• 该框架成功识别出MIRAI恶意软件痕迹中的关键特征，如'ATTACK_VECTOR'变量，提供可操作且可解释的洞察。
• 在大规模神经网络工作负载中，该方法将计算时间从数小时缩短至数秒，实现了可解释AI的实时部署。
• 该方法在多种领域中均具有泛化性与有效性，涵盖图像分类（CIFAR-100）与恶意软件检测（ResNet50）。