2025-11-20T09:19:22.153634

Jet Functors and Weil Algebras in Automatic Differentiation: A Geometric Analysis

Sangha

We present a geometric formulation of automatic differentiation (AD) using jet bundles and Weil algebras. Reverse-mode AD emerges as cotangent-pullback, while Taylor-mode corresponds to evaluation in a Weil algebra. From these principles, we derive concise statements on correctness, stability, and complexity: a functorial identity for reverse-mode, algebraic exactness of higher-order derivatives, and explicit bounds on truncation error. We further show that tensorized Weil algebras permit one-pass computation of all mixed derivatives with cost linear in the algebra dimension, avoiding the combinatorial blow-up of nested JVP/VJP schedules. This framework interprets AD theory through the lens of differential geometry and offers a foundation for developing structure-preserving differentiation methods in deep learning and scientific computing. Code and examples are available at https://git.nilu.no/geometric-ad/jet-weil-ad.

academic

ジェット関手とWeil代数による自動微分：幾何学的分析

基本情報

論文ID: 2510.14342
タイトル: Jet Functors and Weil Algebras in Automatic Differentiation: A Geometric Analysis
著者: Amandip Sangha (ノルウェー、気候環境研究所NILU)
分類: cs.LG math.DG stat.ML
発表日: 2025年10月16日
論文リンク: https://arxiv.org/abs/2510.14342

要約

本論文は、ジェット束(jet bundles)とWeil代数に基づく自動微分(AD)の幾何学的定式化を提案している。逆モードADは余接引き戻し(cotangent-pullback)として表現され、TaylorモードはWeil代数における評価に対応する。これらの原理に基づいて、著者は正確性、安定性、複雑性に関する簡潔な陳述を導出している：逆モードの関手恒等式、高階導数の代数的正確性、および打ち切り誤差の明示的な界。さらに、著者はテンソル化Weil代数により、代数次元に線形な計算コストで、すべての混合導数を一度に計算できることを証明している。これにより、ネストされたJVP/VJPスケジューリングの組み合わせ爆発を回避できる。本フレームワークは微分幾何学の視点からAD理論を解釈し、深層学習および科学計算における構造保存微分法の開発に基礎を提供する。

研究背景と動機

核心問題

自動微分(Automatic Differentiation, AD)は現代機械学習および科学計算の基礎技術であるが、既存のAD理論は統一された幾何学的理論フレームワークを欠いており、以下の問題を引き起こしている：

理論の分散性：逆モードAD(逆伝播)および高階ADの理論的基礎が異なる数学フレームワークに分散している
複雑性の爆発：高階混合導数の計算が組み合わせ複雑性の爆発に直面している
不変性の欠如：既存の方法が座標無関の幾何学的解釈を欠いており、安定性分析に影響を与えている

研究の重要性

本研究は以下の点で重要である：

理論の統一：ADに統一された微分幾何学的理論基礎を提供する
計算効率：高階導数計算の効率問題を解決する
応用展望：深層学習における幾何学的認識手法に理論的支援を提供する

既存方法の限界

従来のAD方法：座標表現に依存し、幾何学的不変性を欠く
高階導数計算：ネストされたJVP/VJP方法が指数級の複雑度を持つ
安定性分析：体系的な誤差伝播理論を欠く

核心貢献

逆伝播の幾何学的理論の確立：逆モードADが余接引き戻し操作と等価であることを証明し、座標無関の表現を提供する
Weil代数フレームワークの提案：Taylorモードを逆モードADをWeil代数における正確な評価として定式化し、代数的正確性を保証する
テンソル化Weil代数法の開発：すべての混合導数の一度の計算を実現し、複雑度が代数次元に線形である
完全な理論分析の提供：正確性証明、安定性界限、複雑度分析を含む

方法の詳細

タスク定義

滑らかな写像 $f: M \to N$ （ここで $M, N$ は滑らかな多様体）とスカラー関数 $\ell: N \to \mathbb{R}$ が与えられたとき、目標は：

合成関数 $\ell \circ f$ の勾配を計算する
$f$ の高階導数を計算する
幾何学的に不変な方法で上記計算を実装する

核心理論フレームワーク

1. 逆モードADの幾何学的表現

定理1（逆伝播を余接引き戻しとして）：滑らかな写像 $f: M \to N$ と $\ell: N \to \mathbb{R}$ に対して、以下が成立する： $\nabla_x(\ell \circ f) = (df_x)^*(d\ell_{f(x)})$

等価的に、ジェット層では： $(j^1f)^*(j^1\ell) = j^1(\ell \circ f)$

この定理は逆伝播を余接空間上の引き戻し操作として再定式化し、以下の幾何学的意義を持つ：

座標無関性：結果は特定の座標系の選択に依存しない
関手性質： $(d(g \circ f)_x)^* = (df_x)^* \circ (dg_{f(x)})^*$ を満たす
自然性：滑らかな再パラメータ化と両立する

2. Weil代数におけるTaylorモード

定理2（Weil模式評価の正確性）： $m^{k+1} = 0$ を満たすWeil代数 $W$ に対して、提升写像 $T_W f: T_W U \to T_W \mathbb{R}^m$ は $f$ の $x$ における全ての $k$ 階導数を打ち切られたTaylor展開の係数として正確に計算する。

Weil代数の構成：

$W = \mathbb{R}[\varepsilon]/(\varepsilon^{k+1})$ またはテンソル積形式
冪零性 $\varepsilon^{k+1} = 0$ が自動的に打ち切りを実現する
代数演算が導数伝播規則に直接対応する

3. テンソル化Weil代数

定理3（テンソル化Weil代数の複雑度）：テンソル化Weil代数を考える： $W \cong \bigotimes_{j=1}^p \mathbb{R}[\varepsilon_j]/(\varepsilon_j^{\rho_j+1}), \quad \dim W = \prod_{j=1}^p (\rho_j + 1)$

$W$ -点 $x_W := x + \sum_{j=1}^p \varepsilon_j v^{(j)}$ における $f$ の単一評価により全ての混合方向導数が得られ、時間複雑度は $O(\dim W \cdot Q)$ である。ここで $Q$ は元のプログラムのスカラー演算数である。

技術的革新点

幾何学的統一性：ADのすべてのモードを微分幾何学フレームワークの下で初めて統一する
代数的正確性：冪零性を通じた打ち切りの代数的正確性を実現し、数値誤差を回避する
線形複雑度：テンソル化法が従来のネスト法の組み合わせ爆発を回避する
逆テープ不要：Weilモードは係数配列のみを保存し、計算グラフ保存を必要としない

実験設定

理論検証

著者は主に理論分析を通じて方法の有効性を検証しており、以下を含む：

正確性検証：関手性質を通じた証明
安定性分析：明示的な誤差界限の提供
複雑度分析：従来の方法との理論的比較

安定性分析

補題1（逆スキャンの後方安定性）：原始演算 $\{\phi_i\}_{i=1}^L$ を持つ直線プログラムに対して、各伴随 $\phi_i^*$ が以下を満たす場合： $\|\phi_i^*(v)\| \leq L_i\|v\|, \quad \|\hat{\phi}_i^*(v) - \phi_i^*(v)\| \leq \delta_i\|\phi_i^*(v)\|$

計算された引き戻しは以下を満たす： $\|\hat{f}^*(\bar{y})\| \leq \left(\prod_{i=1}^L (1+\delta_i)L_i\right)\|\bar{y}\|$

複雑度比較

方法	時間複雑度	空間複雑度	テープ必要
ネストJVP/VJP	$O(\binom{p+k}{k} \cdot Q)$	$O(L)$ (テープ)	はい
テンソル化Weil	$O(\prod_{j=1}^p(\rho_j+1) \cdot Q)$	$O(\dim W)$	いいえ

実験結果

理論結果の検証

係数増長包絡

系1： $f \in C^{k+1}(B_r(x), \mathbb{R}^m)$ であり、その導数が $\|D^\ell f(z)\| \leq M_\ell$ を満たすと仮定すると、Taylor係数は以下を満たす： $\|f_\alpha(x)\| \leq \frac{M_{|\alpha|}}{\alpha!}$