[論文レビュー] tntorch: Tensor Network Learning with PyTorch

Mikhail Usvyatsov, Rafael Ballester‐Ripoll+1 more|arXiv (Cornell University)|Jun 22, 2022

Tensor decomposition and applications被引用数 10

ひとこと要約

tntorch は、CP、Tucker、Tensor Train の複数の低ランクテンソル分解を統合した、PyTorch に統合されたテンソルネットワークライブラリです。単一の 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 要素を含む混合テンソルネットワークアーキテクチャを、1 つの計算グラフ内で組み合わせて使用できること。
• クロス近似、ランクの切り詰め、テンソル代数など、効率的で微分可能な操作を提供し、機械学習およびデータ解析に活用できること。
• テンソル補完、モデル圧縮、ブラックボックスサンプリングを用いたサーヴェイモデルなど、高度なユースケースをサポートする柔軟で拡張可能なフレームワークを提供すること。

提案手法

• すべてのテンソル分解を統合した `Tensor` クラスに封入し、PyTorch と同等の API を提供することで、ディープラーニングワークフローへのシームレスな統合を実現。
• テンソルを低ランクテンソルネットワークとして表現し、各ノードをコアまたは要因として扱い、自動勾配計算を伴うテンソル畳み込みによって演算を実行。
• CP、Tucker、TT 要素が単一のテンソルオブジェクト内で任意の組み合わせで共存・組み合わせ可能となるように、混合フォーマットテンソルネットワークをサポート。
• ブラックボックステンソル関数から重要なインデックスをサンプリングする maxvol アルゴリズムを用いた微分可能クロス近似を実装し、全テンソルを保存せずにメモリ効率の良い学習を可能に。
• 全テンソルを解凍せずに効率的な行列-ベクトル乗算を実現するための特別なクラス `TTMatrix` および `CPMatrix` を提供。
• PyTorch の autograd システムを活用して、グローバル最適化およびリーマン最適化を含む勾配ベースの最適化を可能にし、低ランクテンソル構造の学習を支援。

実験結果

リサーチクエスチョン

• RQ1 CP、Tucker、TT 分解の複数のテンソル分解形式を、微分可能かつ GPU 対応を保ちつつ、1 つの合成可能なインターフェースに統合するにはどうすればよいか？
• RQ2 クロス近似を微分可能にし、大規模またはメモリ上に存在するテンソルの勾配ベース学習を、ディープラーニングパイプラインに統合できるか？
• RQ3 CP-TT-Tucker などのハイブリッド型混合テンソルネットワークアーキテクチャは、機械学習タスクにおける表現力と効率性をどの程度向上できるか？
• RQ4 畳み込み、パディング、ランクの切り詰めといったテンソル代数演算を、微分可能かつスケーラブルな方法で効率的に実装できるか？
• RQ5 PyTorch ネイティブのライブラリが、専用のテンソルライブラリと同等のパフォーマンスと柔軟性を達成できるか、かつ使いやすさと拡張性を維持できるか？

主な発見

• tntorch は、自動微分および GPU 計算を完全にサポートする `Tensor` クラスに、CP、Tucker、TT 分解を統合し、成功裏に実現した。
• maxvol アルゴリズムを用いた微分可能クロス近似を実装し、全テンソルを保存せずにブラックボックス関数から TT テンソルを効率的に学習可能にした。
• CP-TT-Tucker などの混合テンソルネットワークアーキテクチャがネイティブにサポートされており、標準的な PyTorch 演算を用いてエンドツーエンドで構築・最適化可能である。
• 要素ごとの関数、畳み込み、ランクの切り詰めを含む微分可能なテンソル代数演算が、最小限のメモリオーバーヘッドで効率的に計算可能である。
• `TTMatrix` および `CPMatrix` の実装により、テンソルを解凍せずに線形時間オーダー（コア数に比例）で行列-ベクトル乗算を高速に実行できる。
• 感度解析、統計的モーメント、TT 圧縮確率分布からのサンプリングといった高度な機能もサポートしており、確率的モデリングや最適化分野への応用範囲を広げている。