技術要約: Tensor Network Kernel Machines: 機械学習および非線形システム同定のためのJAXフレームワーク

問題提起

高い表現力と計算効率のバランスを両立させた非線形モデルの開発は、機械学習および非線形システム同定において依然として重要な課題である。カーネル法、ガウス過程、深層ニューラルネットワークなどの従来のアプローチは、パラメータ数や訓練サンプル数が増加するにつれて、モデルの複雑さや計算コストが指数関数的に増大するというスケーラビリティの問題に直面することが多い。具体的には、高次元空間における明示的な特徴写像（feature maps）を介して非線形写像を表現しようとすると、パラメータ数が指数関数的に爆発し、高次元問題に対する直接的な最適化が不可能になる。

手法

本論文では、明示的な非線形特徴表現と、コンパクトで低ランクなテンソルネットワーク・パラメトリゼーションを組み合わせることでスケーラビリティに対処するフレームワーク、Tensor Network Kernel Machines (TNKM) を提案する。核となる数学的定式化は、非線形関数 $f(x)$ を高次元特徴空間における線形モデルとして表現するものである：
$$f(x) = \phi(x)^\top w$$
ここで、$\phi(x)$ は非線形特徴写像であり、$w$ はパラメータベクトルである。$w$ に伴う「次元の呪い」を克服するために、TNKMアプローチは2つの主要な戦略を採用している：

1. テンソル積特徴写像: 特徴写像 $\phi(x)$ は、個々の入力変数に作用する1次元写像のテンソル積として構成される。これは積カーネルから自然に導かれるものである。
2. 低ランクテンソル分解: フルパラメータベクトル $w$ をそのまま保存する代わりに、パラメータをCanonical Polyadic (CP) および Tensor-Train (TT) 形式といった低ランクテンソル分解を用いて表現する。これにより、パラメータの複雑さを入力次元に対して指数関数的から線形（またはそれに近いレベル）へと削減する。

著者らは、このフレームワークを実装するために、JAX上に構築されたオープンソースのPythonライブラリである tnkm を開発した。このライブラリは、以下の3つの主要コンポーネントを分離するモジュール式アーキテクチャを採用している：

• 特徴写像 (Feature Maps): 多項式、フーリエ、およびVolterra（メモリベース）の特徴をサポートする。
• テンソルネットワーク・パラメトリゼーション: 専用のモデルクラス（例：CPKM, TTKM）を介してCP分解およびTT分解を実装する。
• 最適化アルゴリズム: 2つの補完的な戦略を提供する：
  • 交互最小二乗法 (Alternating Least Squares: ALS): 多重線形構造を利用して、最小二乗問題に対して高速かつ決定論的な収束を実現する。
  • 勾配ベース最適化: JAXの自動微分機能とOptaxライブラリ（例：Adam, SGD）を活用し、任意の微分可能な損失関数に対する柔軟性を提供する。

主な貢献

本論文は、主に4つの貢献を述べている：

1. 統一されたJAXフレームワーク: 異なる特徴写像、テンソルネットワーク・バックボーン（CPおよびTT）、および最適化戦略を単一のインターフェースの下に統合した、依存関係の少ないオープンソースライブラリ（tnkm）を提供すること。
2. スケーラブルな訓練実装: 構造を考慮したALS最適化を一般的な勾配ベースの手法と併せて統合することで、回帰およびより広範な学習タスクへの適用を可能にしたこと。
3. モジュール式ソフトウェアアーキテクチャ: コンポーネント（特徴写像、モデルクラス、オプティマイザ）を体系的に入れ替えることが可能な設計を行い、非線形システム同定に特化したデータ前処理、シミュレーション、および評価のためのユーティリティを提供すること。
4. 包括的な実験的評価: 非線形回帰およびシステム同定タスクにおける厳格なベンチマークを実施し、効率的な訓練時間とともに競争力のある精度を実証したこと。

実験結果

著者らは、以下の2つの側面からtnkmを評価した。

1. 最適化手法の比較 (Airfoil Self-Noise データセット):

• ALS vs. Adam: ALSは、Kernel Ridge Regression (KRR) のベースライン（0.15）に近い検証用平均二乗誤差（MSE）0.102へと迅速に収束し、わずか1.16秒で達成された。対照的に、Adamはより高い誤差（0.145）に達するまでに、ミニバッチで3.25秒を要し、より多くのエポックを必要とした。
• 結論: ALSは、最小二乗問題としてのTNKMにおいて優れた効率性を示す一方、勾配ベースの手法は非最小二乗目的関数に対する必要な柔軟性を提供する。

2. 非線形システム同定ベンチマーク:
フレームワークは、Coupled Electric Drives、Cascaded Tanks、および Silverbox の3つの標準的なベンチマークでテストされた。

• Coupled Electric Drives: TNKMは、競合手法（GPNARXおよびdynoNet）と同等のRMSEを達成したが、訓練時間は1秒未満であった。これに対し、競合手法は数秒から数分を要した。
• Cascaded Tanks: TNKMは、ブラックボックス手法の中で競争力のある精度（RMSE 0.347）を提供し、SUBNETやPNARXといった手法を訓練速度において大幅に上回った（TNKMは1秒未満に対し、他は1〜60分）。
• Silverbox: TNKMは、主要なブラックボックス手法（GPNARX, SUBNET）と同等の性能を達成したが、訓練時間は1〜60秒であった。物理情報に基づく手法（NL-LFR）の精度には及ばなかったものの、長期予測における安定性を実証した。

意義と主張

本論文は、tnkmがテンソルネットワークに基づく学習手法の再現可能な開発と適用を促進すると主張している。特徴写像、テンソル・パラメトリゼーション、および最適化アルゴリズムを切り離すことで、本ライブラリは、基礎となるワークフローを変更することなく、TNKMの構成を体系的に探索することを可能にする。

著者らは、TNKMを古典的なカーネル法と深層ニューラルネットワークの中間的なアプローチとして位置づけている。カーネル法と同様に、TNKMは明示的な特徴表現を使用するため、制御可能性が高く、事前知識を組み込むことが容易である。しかし、フルパラメータベクトルをコンパクトなテンソルネットワークに置き換えることでスケーラビリティを確保している。ニューラルネットワークとは異なり、特徴をデータから直接学習するのではなく、著者らが認めるように、手動による特徴選択が必要となるという制限がある。

本研究の意義は、理論的なテンソルネットワーク手法と現実世界の非線形システム同定との間の溝を埋める、実用的で拡張可能なソフトウェア基盤を提供した点にある。これにより、これらのモデルが高い予測精度を維持しながら、低い計算コストを実現できることを示した。