[논문 리뷰] tntorch: Tensor Network Learning with PyTorch

Mikhail Usvyatsov, Rafael Ballester‐Ripoll외 1명|arXiv (Cornell University)|2022. 06. 22.

Tensor decomposition and applications인용 수 10

한 줄 요약

tntorch는 PyTorch에 통합된 텐서 네트워크 라이브러리로, CP, Tucker, Tensor Train 등 여러 가지 저질서 텐서 분해를 단일하고 PyTorch 유사한 API로 통합합니다. 자동 미분, GPU 가속, 원활한 미분 가능한 연산을 지원하여 고차원 텐서를 효율적으로 학습, 조작, 압축할 수 있으며, 혼합 형식의 네트워크와 고급 기능인 크로스-근사 및 미분 가능한 텐서 대수를 지원합니다.

ABSTRACT

We present tntorch, a tensor learning framework that supports multiple decompositions (including Candecomp/Parafac, Tucker, and Tensor Train) under a unified interface. With our library, the user can learn and handle low-rank tensors with automatic differentiation, seamless GPU support, and the convenience of PyTorch's API. Besides decomposition algorithms, tntorch implements differentiable tensor algebra, rank truncation, cross-approximation, batch processing, comprehensive tensor arithmetics, and more.

연구 동기 및 목표

• CP, Tucker, TT 등의 다양한 텐서 분해 형식을 상호 교환 가능한 인터페이스로 통합하여 기저 표현 방식의 차이를 추상화합니다.
• PyTorch의 자동 미분 및 GPU 가속과 완전히 호환되는 환경을 제공하여 종단 간 미분 가능한 텐서 학습을 실현합니다.
• 단일 계산 그래프 내에서 CP, Tucker, TT 구성 요소를 혼합한 텐서 네트워크 아키텍처를 지원합니다.
• 크로스-근사, 랭크 절단, 텐서 대수와 같은 효율적인 미분 가능한 연산을 제공하여 기계 학습 및 데이터 분석에 활용합니다.
• 텐서 완성, 모델 압축, 블랙박스 샘플링을 통한 서rogate 모델링과 같은 고급 사용 사례를 지원하는 유연하고 확장 가능한 프레임워크를 제공합니다.

제안 방법

• 모든 텐서 분해를 단일한 `Tensor` 클래스로 봉인하여 PyTorch 호환 인터페이스를 제공함으로써 딥러닝 워크플로우와 원활하게 통합됩니다.
• 각 노드가 코어 또는 요소로 구성된 저질서 텐서 네트워크로 텐서를 표현하며, 자동 기울기 계산이 가능한 텐서 결합 연산을 수행합니다.
• CP, Tucker, TT 구성 요소가 단일 텐서 객체 내에서 서로 다른 조합으로 공존하고 조합될 수 있도록 혼합 형식 텐서 네트워크를 지원합니다.
• 블랙박스 텐서 함수로부터 핵심 색인을 추출하기 위해 maxvol 알고리즘을 사용한 미분 가능한 크로스-근사를 구현하여 전체 텐서 저장 없이 메모리 효율적인 학습을 가능하게 합니다.
• 전체 텐서를 해체하지 않고도 빠른 행렬-벡터 곱셈을 수행할 수 있도록 `TTMatrix` 및 `CPMatrix`와 같은 전용 클래스를 제공합니다.
• PyTorch의 autograd 시스템을 활용하여 기울기 기반 최적화, 특히 전역 최적화 및 리만 기반 최적화를 통해 저질서 텐서 구조 학습을 지원합니다.

실험 결과

연구 질문

• RQ1 CP, Tucker, TT 등의 다양한 텐서 분해 형식을 상호 작용 가능한 인터페이스로 통합할 수 있는 방법은 무엇인가요? 이 인터페이스는 미분 가능성과 GPU 지원을 유지할 수 있나요?
• RQ2 크로스-근사는 어떻게 미분 가능하게 만들 수 있으며, 대규모 또는 메모리에 올라온 텐서의 기울기 기반 학습을 위한 딥러닝 파이프라인에 통합할 수 있나요?
• RQ3 CP-TT-Tucker 하이브리드와 같은 혼합 텐서 네트워크 아키텍처는 기계 학습 작업에서 표현력과 효율성을 얼마나 향상시킬 수 있나요?
• RQ4 텐서 대수 연산, 예를 들어 컨볼루션, 패딩, 랭크 절단은 어떻게 효율적이고 미분 가능하며 확장 가능한 방식으로 구현할 수 있나요?
• RQ5 PyTorch 네이티브 라이브러리가 전용 텐서 라이브러리와 비교해 성능와 유연성 면에서 유사한 성능을 달성하면서도 사용 용이성과 확장성은 유지할 수 있나요?

주요 결과

• tntorch는 자동 미분 및 GPU 계산을 완전히 지원하는 단일 `Tensor` 클래스를 통해 CP, Tucker, TT 분해를 성공적으로 통합합니다.
• maxvol 알고리즘을 통한 미분 가능한 크로스-근사를 통해 블랙박스 함수에서 전체 텐서를 저장하지 않고도 TT 텐서의 효율적 학습이 가능합니다.
• 혼합 텐서 네트워크 아키텍처(예: CP-TT-Tucker)가 내장되어 있으며, 표준 PyTorch 연산을 사용해 종단 간 구성 및 최적화가 가능합니다.
• 요소별 함수, 컨볼루션, 랭크 절단과 같은 미분 가능한 텐서 대수 연산은 최소한의 메모리 오버헤드로 효율적으로 계산할 수 있습니다.
• `TTMatrix` 및 `CPMatrix`의 구현을 통해 텐서를 해체하지 않고도 선형 시간 복잡도(핵심 수에 비례)로 빠른 행렬-벡터 곱셈을 수행할 수 있습니다.
• 감도 분석, 통계적 모멘트, TT 압축 확률 분포에서의 샘플링과 같은 고급 기능을 지원하여 확률 모델링 및 최적화 분야의 응용 범위를 넓힙니다.