[논문 리뷰] projUNN: efficient method for training deep networks with unitary matrices

Bobak T. Kiani, Randall Balestriero외 2명|arXiv (Cornell University)|2022. 03. 10.

Domain Adaptation and Few-Shot Learning인용 수 11

한 줄 요약

이 논문은 단위행렬을 사용한 딥 네트워크 훈련을 위한 효율적인 방법인 projUNN을 제안한다. 이 방법은 단위성을 유지하기 위해 투영 또는 운반을 통해 질량-𝑘 업데이트를 사용하며, 근사적으로 최적의 $O(kN^2)$ 훈련 복잡도를 달성한다. 이는 기존 최고 수준의 단위행렬 RNN을 능가하거나 동등하게 유지하면서, 정규화된 컨volutional 네트워크의 안정적이고 확장 가능한 훈련을 가능하게 한다.

ABSTRACT

In learning with recurrent or very deep feed-forward networks, employing unitary matrices in each layer can be very effective at maintaining long-range stability. However, restricting network parameters to be unitary typically comes at the cost of expensive parameterizations or increased training runtime. We propose instead an efficient method based on rank-$k$ updates -- or their rank-$k$ approximation -- that maintains performance at a nearly optimal training runtime. We introduce two variants of this method, named Direct (projUNN-D) and Tangent (projUNN-T) projected Unitary Neural Networks, that can parameterize full $N$-dimensional unitary or orthogonal matrices with a training runtime scaling as $O(kN^2)$. Our method either projects low-rank gradients onto the closest unitary matrix (projUNN-T) or transports unitary matrices in the direction of the low-rank gradient (projUNN-D). Even in the fastest setting ($k=1$), projUNN is able to train a model's unitary parameters to reach comparable performances against baseline implementations. In recurrent neural network settings, projUNN closely matches or exceeds benchmarked results from prior unitary neural networks. Finally, we preliminarily explore projUNN in training orthogonal convolutional neural networks, which are currently unable to outperform state of the art models but can potentially enhance stability and robustness at large depth.

연구 동기 및 목표

• 단위행렬을 사용한 딥 또는 순환 네트워크 훈련의 높은 계산 비용 문제를 해결하기 위해, 기존에 단위행렬 업데이트에 $O(N^3)$ 연산이 필요한 문제를 해결한다.
• 특히 큰 $N$에 대해 높은 훈련 런타임을 유지하면서도 $N\times N$ 단위행렬의 전체 매개변수화를 가능하게 한다.
• 비싼 SVD나 행렬 역행렬 계산 단계 없이 단위성을 유지함으로써 딥 네트워크의 안정성과 확장성을 향상시킨다.
• 특히 깊은 아키텍처에서 단위행렬/정규화된 행렬의 적용 가능성을 확장하기 위해 컨볼루션 네트워크에 적용 가능하도록 한다.

제안 방법

• projUNN는 단위행렬 업데이트의 계산 비용을 $O(N^3)$에서 $O(kN^2)$로 줄이기 위해 질량-𝑘 업데이트를 사용하여 기울기의 근사치를 계산한다.
• 직접 변환 버전인 projUNN-D는 낮은 랭크 업데이트 부분공간의 고유분해를 통해 업데이트된 행렬을 가장 가까운 단위행렬로 투영한다.
• 접선 변환 버전인 projUNN-T는 지수 매핑 업데이트를 통해 단위행렬을 낮은 랭크 기울기 방향으로 운반한다.
• 두 버전 모두 전체 단위행렬 다양체를 매개변수화하고 각 기울기 업데이트 후 단위성을 강제로 유지함으로써 수치적 안정성을 확보한다.
• 컨볼루션 레이어의 경우, FFT를 사용하여 푸리에 도메인에서 정규화된 컨볼루션을 수행함으로써 허미트 대칭성을 유지하고 효율적인 계산을 가능하게 한다.
• 단위행렬/정규화된 매개변수를 위해 특수한 최적화 알고리즘을 사용하며, 일반적으로 표준 매개변수보다 1/10 또는 1/20 수준의 학습률을 설정한다.

실험 결과

연구 질문

• RQ1 성능을 저하시키지 않고도 단위행렬을 딥 네트워크에서 효율적으로 훈련하기 위해 저랭크 기울기 근사치를 사용할 수 있는가?
• RQ2 제안된 방법은 단위행렬 다양체의 전체 매개변수화를 유지하면서도 근사적으로 최적의 훈련 복잡도를 달성하는가?
• RQ3 장거리 시퀀스 작업에서 기존의 단위행렬 RNN과 비교해 projUNN의 정확도와 훈련 속도는 어떠한가?
• RQ4 projUNN는 정규화된 컨볼루션 네트워크로 효과적으로 확장되어 깊은 아키텍처에서의 안정성을 향상시킬 수 있는가?
• RQ5 Haar 랜덤 또는 항등행렬 초기화와 비교해 어떤 초기화 전략이 projUNN에서 더 나은 수렴성과 안정성을 제공하는가?

주요 결과

• projUNN는 기울기 단계당 $O(kN^2)$의 훈련 복잡도를 달성하여 이전의 리 대수 기반 방법의 $O(N^3)$ 비용을 크게 감소시킨다.
• 조금이라도 $k=1$일 때도 projUNN는 벤치마크 시퀀스 학습 작업에서 최고 수준의 단위행렬 RNN의 성능을 따라하거나 능가한다.
• 직접 변환 버전인 projUNN-D에 비해 접선 변환 버전인 projUNN-T는 많은 훈련 단계 동안 더 뛰어난 수치적 안정성을 보였다.
• projUNN는 정규화된 컨볼루션 네트워크의 효율적 훈련을 가능하게 하여 매우 깊은 모델에서의 안정성 향상에 기여한다.
• 경험적으로, $2\times2$ 비대칭 블록을 사용한 초기화 방식(헤나프 스타일)이 항등행렬 또는 Haar 랜덤 초기화보다 성능이 뛰어나다.
• 매 $O(N)$ 단계마다 간헐적으로 전체 투영을 수행함으로써 단위성을 유지하면서도 평균 $O(kN^2)$ 런타임이 저하되지 않는다.