[논문 리뷰] Fast Differentiable Matrix Square Root

Song Yue, Nicu Sebe외 1명|arXiv (Cornell University)|2022. 01. 21.

Advanced Vision and Imaging인용 수 5

한 줄 요약

이 논문은 효율적인 전방 계산을 위한 두 가지 빠른 미분 가능한 행렬 제곱근 방법—행렬 테일러 다항식(MTP)과 행렬 페드에 근사법(MPA)—을 제안한다. 이는 백프로파게이션에 대해 반복적인 리아푸노프 해법기를 통합한 것으로, SVD 및 뉴턴-슐츠 반복보다 빠른 성능을 달성하면서도 시각 작업에서 경쟁적인 성능을 유지한다.

ABSTRACT

Computing the matrix square root or its inverse in a differentiable manner is important in a variety of computer vision tasks. Previous methods either adopt the Singular Value Decomposition (SVD) to explicitly factorize the matrix or use the Newton-Schulz iteration (NS iteration) to derive the approximate solution. However, both methods are not computationally efficient enough in either the forward pass or in the backward pass. In this paper, we propose two more efficient variants to compute the differentiable matrix square root. For the forward propagation, one method is to use Matrix Taylor Polynomial (MTP), and the other method is to use Matrix Padé Approximants (MPA). The backward gradient is computed by iteratively solving the continuous-time Lyapunov equation using the matrix sign function. Both methods yield considerable speed-up compared with the SVD or the Newton-Schulz iteration. Experimental results on the de-correlated batch normalization and second-order vision transformer demonstrate that our methods can also achieve competitive and even slightly better performances. The code is available at \href{https://github.com/KingJamesSong/FastDifferentiableMatSqrt}{https://github.com/KingJamesSong/FastDifferentiableMatSqrt}.

연구 동기 및 목표

• 딥 러닝 프레임워크에서 특히 배치 처리된 행렬 제곱근 연산에 대해 SVD 및 뉴턴-슐츠 반복의 계산 비효율성을 해결한다.
• 고유값 차이로 인한 고유분해의 불안정성과 높은 비용으로 인해 백프로파게이션에서 발생하는 문제를 해결한다.
• GPU에 적합하고, 빠르며 수치적으로 안정적인, 미분 가능한 행렬 제곱근 방법을 개발한다.
• 시각 응용 분야(예: 공분산 풀링 및 정규화 레이어)에서 행렬 제곱근의 효율적 사용을 가능하게 한다.

제안 방법

• 전방 계산에서 SVD 또는 반복적 방법을 대체하기 위해 행렬 테일러 다항식(MTP)과 행렬 페드에 근사법(MPA)을 사용해 행렬 제곱근을 근사한다.
• 백프로파게이션 중에 기저가 행렬 부호 함수인 반복적 리아푸노프 해법기를 적용해 기울기를 효율적으로 계산한다.
• 역전파를 연속 시간 리아푸노프 방정식을 푸는 것으로 공식화하여 명시적 고유분해를 피한다.
• 행렬 곱셈이 GPU에서 잘 최적화되어 있으므로, 최소한의 계산 오버헤드로 딥 러닝 프레임워크에 통합한다.
• 반복적 구성요소의 수렴 안정성을 향상시키기 위해 입력 행렬을 사전 정규화한다.
• 대규모 모델에서 특히 효율성을 높이기 위해 혼합 정밀도 학습 호환성을 활용한다.

실험 결과

연구 질문

• RQ1 MTP와 MPA는 딥 러닝에서 행렬 제곱근 계산에 대해 SVD의 더 빠르고, 미분 가능한 대안이 될 수 있는가?
• RQ2 제안된 리아푸노프 기반 역전파 해법기는 기울기 계산 속도 측면에서 전통적인 SVD 및 뉴턴-슐츠 반복보다 뛰어나게 성능을 내는가?
• RQ3 실제 시각 벤치마크(예: ZCA 화이트닝 및 2차 시각 트랜스포머)에서 제안된 방법들은 속도와 정확도 측면에서 어떻게 비교되는가?
• RQ4 대규모 모델에서 학습 시간과 메모리 사용량을 줄이면서도 성능을 유지하거나 향상시킬 수 있는가?

주요 결과

• So-ViT-14 모델에서 우리의 MPA-Lya 방법은 뉴턴-슐츠 반복보다 3.19배 빠르고, SVD-Padé보다 25.63배 빠르다.
• MTP-Lya 방법은 행렬 제곱근 계산당 2.36ms로 가장 빠른 전방 계산 성능를 기록하며, SVD 및 NS 반복을 능가한다.
• CIFAR-100에서 MPA-Lya는 2.52ms의 추론 시간과 함께 19.55% ± 0.20의 검증 오차를 기록하여, 속도와 정확도에서 SVD 및 NS 반복을 모두 앞선다.
• 2차 시각 트랜스포머에서 MPA-Lya는 254 에포크 만에 수렴하며, SVD 기반 방법보다 적은 에포크를 소요하지만, 상위 1위 정확도는 동일하거나 略적으로 초월한다.
• MTP-Lya 방법은 ResNet-18에서는 SVD 및 NS와 유사한 성능를 기록하지만, ResNet-50에서는 약간의 성능 저하가 발생한 것으로 보이며, 이는 근사 오차 누적이 원인일 수 있다.
• 리아푸노프 기반 역전파 해법기는 NS 반복에 비해 행렬 곱셈 횟수를 줄여 더 빠른 백프로파게이션을 가능하게 한다.