Implicit differentiation of tensor network algorithms

技術要約：テンソルネットワーク・アルゴリズムの陰関数微分

問題提起
基底状態計算のための投影絡み合い対状態（PEPS）の変分最適化は、局所的な変分パラメータに対するエネルギー勾配の正確な計算に大きく依存している。現在の最先端のアプローチでは、自動微分（AD）を用いてこれらの勾配を計算している。しかし、この手法は以下の3つの重大な課題に直面している：

1. 効率性： 標準的なリバースモードADは、収縮アルゴリズム（例：CTMRGや境界MPS）のすべての反復過程を通じてバックプロパゲーションを行う必要がある。固定点定式化を用いて完全な計算グラフの保存を回避する場合でも、反復的なサブルーチン（疎行列固有値ソルバーやSVDなど）の存在により、勾配計算において入れ子状の反復手順が生じ、計算コストが高くなる。
2. 安定性： 収縮アルゴリズムで使用される原始的なサブルーチン（例：切り捨てられた固有値分解やSVD）を微分することは、数値的不安定性を導入する。これらの不安定性は、固有値や特異値が近接または縮退している場合に発生し、微分値が発散する原因となる。極めて重要な点は、これらの不安定性は微分手順による人工的なものであり、物理的なエネルギー勾配自体は、基礎となる変動に対して安定しているということである。
3. 実装の複雑さ： PEPS最適化にADを適用するには、収縮アルゴリズム内のすべての原始操作がADエンジンと互換性を持つ必要がある。非ネイティブなサブルーチンに対してカスタムの微分規則を定義することは労力が大きく、柔軟性を制限する。

手法
著者らは、これらの課題に対処するため、**陰関数微分（implicit differentiation）**に基づくフレームワークを提案している。この手法は、収縮アルゴリズムの反復ステップを微分する代わりに、収束した収縮環境を単一の代数的な特性方程式の根として再定式化するものである。

本手法の核心となるステップは以下の通りである：

1. 代数的特徴付け： 収束した環境（例：CTMRGにおけるコーナーテンソルとエッジテンソル、または境界MPSテンソル）は、反復更新ルール $x = f(x, p)$ ではなく、一連の特性方程式 $F(y, p) = 0$ によって特徴付けられる。ここで、$y$ は環境変数（中間変数としての等長テンソルや特異値を含む可能性がある）の集合を表し、これらは変分パラメータ $p$ に暗黙的に依存する。
2. 陰関数微分： $F(y, p) = 0$ に対して陰関数定理を適用することで、勾配計算は単一の入れ子構造を持たない線形方程式系の求解へと簡略化される：
   $$\tilde{y} \partial_y F = -\tilde{y} \partial_p F$$
   ここで、$\tilde{y}$ は随伴変数である。これにより、反復的なサブルーチンを介したバックプロパゲーションが排除される。
3. 安定なパラメータ化： 本手法の鍵となる革新は、変数集合 $y$ の選択と特性方程式のパラメータ化にある。著者らは、収縮環境に内在するゲージ対称性を利用することで、数値的不安定性の原因となる自由度（例：縮退部分空間内での変動）を明示的に除去するパラメータ化を構築している。例えば、CTMRGにおいては、等長テンソルの変動をその零空間内に制限することで、切り捨てられた固有値分解の微分に伴う発散項を回避している。
4. モジュール化された実装： このアプローチは、収縮アルゴリズムと勾配計算を分離する。フォワードパスは変更されず（標準的な収縮ルーチンを使用）、リバースパスでは代数的な特性方程式の微分のみが必要となる。これらの方程式は単純なテンソル縮約で構成されている。

主な貢献
本論文は、主要な3つのPEPS収縮スキームに対する特性方程式の明示的な定式化を提供している：

• $C_{4v}$ 対称性を持つCTMRG： 著者らは、コーナーテンソル、エッジテンソル、およびパラメータ化された等長テンソルを含む特性方程式のセットを導出している。この定式化により、切り捨てられた固有値分解（eigh）の微分が不要となり、縮退した固有値に関連する不安定性が解消される。
• 境界MPS (VUMPS)： 上部および下部の境界MPS固有ベクトルと混合環境に関する特性方程式が定式化されている。このアプローチは、勾配計算中の境界MPSテンソルのゲージ固定を回避し、支配的な固有値ソルバーの微分から生じる入れ子状の線形問題を排除する。
• 非対称CTMRG (対称性なし)： 本手法は、対称性のない一般的な系にも拡張されている。著者らは、ゲージ不変性と、切り捨てられたSVDの微分時における安定性を確保するために、修正されたコーナーおよびエッジテンソル、および特異値に対するカスタムの「平方根」変数を導入している。

結果
著者らは、PEPSKit.jlおよびMPSKit.jlパッケージを用い、陰関数微分を固定点微分およびストレートなリバースモードADと比較してベンチマークを行った：

• 精度： 陰関数微分によって計算された勾配は、線形ソルバーの許容誤差範囲内で、参照用の有限差分値およびブラックボックスのリバースモードADと一致しており、これは中程度の収束状態の環境においても同様であった。
• 性能：
  • $C_{4v}$ CTMR_Gおよび境界MPSにおいて、陰関数的アプローチは固定点微分を一貫して上回り、ボンド次元 ($D$) および環境次元 ($\chi$) が大きくなるにつれて加速が増す。
  • 非自明なユニットセルを持つ非対称CTMRGにおいて、陰関数法は優れた漸近的スケーリングを示す。低次元では実装のオーバーヘッドにより固定点微分の方が高速な場合があるが、疎行列分解が必要となる高次元においては、陰関数法が著しく効率的になる。
• スケーリング： 多くの場合、陰関数微分を用いた勾配計算のコストは、収縮コストと比較して劣位のステップとなる。一方、固定点微分はしばしば支配的なボトルネックであり続ける。

意義と主張
本論文は、このフレームワークがテンソルネットワーク・プログラミングにおけるパラダイムシフトを象徴すると主張している。アルゴリズムのサブルーチンではなく、代数的な最適条件を微分することで、本手法は以下を実現する：

• 実装の簡素化： 収縮アルゴリズム内のすべてのサブルーチンが微分可能である必要がなくなり、フォワードパスにおいて標準的で高度に最適化された線形代数ルーチンを使用することが可能になる。
• 安定性の向上： ゲージ対称性を利用して安定なパラメータ化を構築することにより、サブルーチンの微分から生じる数値的不安定性を自然に除去する。
• 効率性の向上： 線形問題をアンネスト（非入れ子化）することで、勾配評価の計算複雑性を低減し、システムサイズに対するスケーリングを改善する。

著者らは、本手法を3つの特定の収縮スキームで実証したが、コアとなるアルゴリズム（アルゴリズム1および2）は一般的であり、解が微分可能な最適条件によって定義できるあらゆる最適化問題に適用可能であることを強調している。彼らは、すべてのPEPSの課題を解決したと主張しているのではなく、このアプローチによって、安定した勾配ベースのPEPS最適化が大幅に実用的かつ堅牢になったと主張している。