Variational inference and density estimation with non-negative tensor of hierarchical tucker format

技術要約：非負階層的タッカー形式を用いた変分推論および密度推定

問題設定
本論文は、高次元の離散確率分布（テンソル $P: [n]^d \to \mathbb{R}_{\geq 0}$ として表現される）を、扱いやすい低ランク形式へと圧縮するという課題に取り組んでいる。このようなテンソルを直接保存するには $O(n^d)$ の複雑さが必要となり、これは中程度の次元 $d$（例：数十の変数）においてさえ実行不可能となる。目標は、$P$ を効率的な保存、サンプリング、および尤度評価が可能なパラメトリックな族へと近似することである。本研究では、以下の2つのアクセスモデルを検討している：

1. 変分推論 (Variational Inference, VI): 定数を除いて $P$ の未正規化値 $P(i_1, \dots, i_d)$ を提供するオラクルを介したアクセス。
2. 密度推定 (Density Estimation, DE): 分布から抽出されたサンプル集合を介した $P$ へのアクセス。

テンソルネットワークモデル（Tensor-Train (TT) 形式など）は、「エリア則（area law）」を満たす分布には効果的であるが、エリア則が成立しない2D格子モデルや複雑な相互作用構造を持つモデルでは苦戦する。さらに、標準的なテンソルネットワークは、確率分布に不可欠な非負性を本質的に保証しておらず、KLダイバージェンスの計算やサンプリングといったダウンストリームのタスクにおける「符号問題（sign problems）」を回避できない。

手法
著者らは、$P$ を 非負階層的タッカー (Non-Negative Hierarchical Tucker, NHT) 形式へと圧縮するための、2段階のエンドツーエンド・アルゴリズムを提案している。

• ステージ 1: 階層的タッカー補間 (参照構成)

  • 目的: 正負の符号を持つ階層的タッカー (Hierarchical Tucker, HT) 近似 $\tilde{P} \approx P$ を構築する。
  • VI の場合: 補間分解（CUR分解のような手法）に基づく既存の補間法を用いて $P$ のエントリーをクエリし、$\tilde{P}$ を再構成する。
  • DE の場合: サンプルから $\tilde{P}$ を構築するための既存の密度推定アルゴリズムを使用する。
  • 出力: 正負の成分を持つ参照テンソル $\tilde{P}$。このステージでは非負性は強制されない。

• ステージ 2: 非負階層的タッカー・フィッティング

  • 目的: 正のテンソル因子 $G = \{G_q\}$ によってパラメータ化された非負アンサンブル $P_G$ を、符号付きの参照 $\tilde{P}$ に適合させる。
  • 最適化: 問題は、$P_G$ と $\tilde{P}$ の間の二乗フロベニウスノルムと、非負性を強制するためのログ障壁ペナルティからなる損失関数 $\ell(G)$ を最小化するように定式化される：
    $$\ell(G) = \|P_G - \tilde{P}\|_F^2 + \sum_q \mu_q \phi_q(G_q)$$
    ここで $\phi_q$ は $G_q$ の各要素に対するログ障壁である。
  • アルゴリズム: 交互最小化アプローチが採用されている。各テンソル成分 $G_q$ に対して、アルゴリズムは二次のニュートンステップを用いて変数を更新する。
    • 勾配/ヘシアン: グラム行列のクロネッカー積を利用した、バイナリツリー構造上のメッセージパッシングを介して計算される。
    • 効率性: 密なヘシアンの逆行列計算に伴う $O(\rho^9)$ のコストを避けるため、中間ノードには共役勾配法 (CG) を、リーフノードにはデカップルされた解法を用いる。
  • ウォーム初期化: 極めて重要なコンポーネントは、「ゲージ・ウォーム初期化（gauge warm initialization）」戦略（アルゴリズム 2.3）である。符号付きの参照を単純に閾値処理するのではなく、この手法はテンソルネットワークにおけるゲージの自由度を活用する。リーフからルートへのデカップルされたスイープを行い、参照成分 $F_q$ に対して「leave-one-out」方式で最もよく適合する直交ゲージ行列 $Q$ と正の成分 $G_q$ を見つけ出す。これにより、収束性が大幅に向上し、局所解を回避できる。

主な貢献

1. 高次元分布のための NHT 形式: 本論文は、Tensor-Train (TT) 形式よりも（2D格子のような）構造的制約に適した、非負階層的タッカー (NHT) 形式の使用を導入している。これにより、確率分布に必要な非負性を維持している。
2. 二次適合手続き: 著者らは、NHT フィッティング・ステージのための専用の二次最適化手法を設計した。これには、高次元の設定における局所解への到達の困難さを解決する、新しいウォーム初期化戦略が含まれており、標準的な初期化と比較して最終的な精度を桁違いに改善することが多い。
3. 計算量: 提案された手法の両ステージは $O(d)$ の計算量を実現しており、高次元設定へのスケーラビリティを備えている。
4. エンドツーエンド・アルゴリズム: 本研究は、変分推論と密度推定の両方のシナリオに対する完全なパイプラインを提供し、符号付きの参照の構築と、物理的に妥当な非負モデルの適合との間の溝を埋めている。

実験結果
数値実験は、両方のアクセスモデルに対して行われた：

• 変分推論: 2D ギンツブルグ・ランダウ・モデル ($d=32$) およびヘビーテイル・モデル ($d=30$) でテストされた。適応的障壁プリコンディショニング付きニュートン反復は、マルチプリカティブ・アップデート（乗法的更新）のベンチマークを大幅に上回り、同一のウォーククロック時間内でより低い相対二乗フロベニウス損失を達成した。
• 密度推定: 2D イジング・モデル ($d=32$) および 2D 横磁場イジング・モデル ($d=64$) でテストされた。提案手法は真の分布を正確に近似し、負の対数尤度 (NLL) は厳密な値に近い値を示した。
• アブレーション研究:
  • NHT vs. NTT: 2D ギンツブルグ・ランダウ・モデルにおいて、NHT アンサンブルは、NTT アンサンブルに2倍のパラメータを与えた場合でも、より優れた適合精度を達成した。
  • 初期化: 提案されたゲージ・ウォーム初期化（アルゴリズム 2.3）は、乗法的更新アルゴリズム（アルゴリズム 3.1）による初期化よりも、一貫して優れた適合結果をもたらした。

意義と主張
本論文は、提案された NHT 手法が、構造的制約（例：2D 格子）により従来のテンソル・トレイン・モデルが失敗する高次元離刻分布の圧縮に対して、堅牢なソリューションを提供することを主張している。アンサンブルを通じて非負性を保証し、特殊化されたウォームスタートを備えた洗練された二次適合手続きを利用することで、本手法は、制約のないテンソルネットワークに関連する符号問題を伴うことなく、正確な密度推定および変分推論を可能にする。著者らは、本手法が完全なバイナリツリー上で示されているものの、コードは不完全なツリーもサポートしており、今後の研究では、関数的階層的テンソル・アンサンブルへの拡張や、モーメント保存制約の組み込みが可能であると述べている。