[논문 리뷰] Hardware Acceleration of Explainable Machine Learning using Tensor Processing Units

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

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.

연구 동기 및 목표

• 실시간 배포를 제한하는 설명 가능한 기계학습(XM)에서의 느린 해석 처리라는 핵심적 병목 현상을 해결하기 위해.
• 반복적이고 시간이 오래 걸리는 최적화 문제에 의존하는 기존 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)를 포함한 다양한 도메인에서 일반화 가능하고 효과적인 것으로 입증되었다.