📂機械学習

論文レビュー：コルモゴロフ・アーノルドニューラルネットワーク（KAN）

概要と要約

Kolmogorov–Arnold Networks（KAN）は、その名の通りコルモゴロフ-アルノルト表現定理^{Kolmogorov–Arnold representation theorem}にインスパイアされて開発されたニューラルネットワークであり、関連するアイデア自体は数十年前から議論されていましたが、ジローシ^Girosiとポッジョ^Poggioらが1989年に「Representation Properties of Networks: Kolmogorov’s Theorem Is Irrelevant」という論文で指摘したように、当該定理と大きな関連性はありません¹。ただし、多くの有名なニューラルネットワークがネットワークの構築方法に従って直感的なネーミングをする中、あえて数学者の名前を取ったのは、それだけ理論的な基盤がしっかりしていることを強調したかったのだと思います。2024年4月にarXivに投稿され²、この投稿を書いている2024年9月27日現在、すでに200回近く引用されているので、所期の目的は達成したのではないかと思います。

KANのコンセプトは確固たるものであり、現代の人工知能といえば絶対的な主流であるMLP^{Multi-Layer Perceptron}と二大勢力を形成し、従来のMLPやディープラーニングがブラックボックス手法として持つ限界を突破しようというものです。性能自体は良いが、その原理を解明しないまま受け入れるのは困難であり、例えば金融や医療、宇宙、軍需などの分野にも進出するためには必ず研究が行われなければならない部分であり、これまでのところかなり斬新なアプローチを示しています。

ただし、モデルのパフォーマンスを含めて学習速度などに関しては、どうしても誇張がかなりあり、肝心な数式化をやや大雑把に解決した点は確かに警戒すべきです。

0. アブストラクト

KANはコルモゴロフ-アルノルト表現定理にインスパイアされており、従来のMLPが既に定められた活性化関数を用いてノードで重みを更新するのとは異なり、KANは活性化関数自体が学習を通じて形を変え、リンクで重みを更新するという戦略を取ります。KANを貫くテーマは解釈可能性^{interpretability}であり、著者らが提案するいくつかの事例では数学者や物理学者にとって有用な共同研究者となりうるとし、MLPに対する有望な代替手段になると主張しています。

1. はじめに

MLPの最大の問題点である解釈可能性から指摘します。MLPがシベンコの定理を理論的根拠として発展したように、KANはコルモゴロフ-アルノルト表現定理を基盤としています。MLPがデータの線形結合に非線形関数を適用し、脳を模倣するコンセプトを持ち、その計算がどれほど複雑で膨大であっても、GPUが実際にそのようなタスクを行列演算で解き始め、大きな発展を遂げました。一方、KANは基本設定からして1次元関数、より具体的にはスプラインを通じてモデルが学習されることを望みます。誰もがこの説明を聞けば自然に「それでは計算に時間がかかりすぎるのではないか？」そして「GPUを使えないということではないか？」と思うでしょうが、幸いなことに通常はMLPに比べてネットワークのサイズが非常に小さくても性能を出せると述べています。

次に、従来もニューラルネットワークを構築するためにコルモゴロフ-アルノルト表現定理を試そうとする試みはありましたが、その定理で明示する条件をあまりに正確に守ろうとしたために現代的な手法を使いにくかったという説明とともに、著者らはそのような構造を任意に変えて様々な試みをしたと述べています。ある意味、コルモゴロフ-アルノルト表現定理と決別したことを告白する場面とも言えますが、実際にも次の段落でKANはスプラインとMLPの組み合わせに過ぎないと認めています。

著者らはアブストラクトから一貫してKANが次元の呪い^{curse of dimension}を克服したという主張を繰り返していますが、正直なところその部分はあまり共感できません。KANの潜在性を示すために様々な問題に挑戦し、依然として数学者や科学者に役立つとアピールしています。正直に言えば、現時点では工学界や産業界の主流を奪うのは難しく、ニューラルネットワークのネーミングから一貫して基礎科学に集中していると見ればよいと思います。研究に使用されたコードはGitHubリポジトリ https://github.com/KindXiaoming/pykan で確認できます。

2. Kolmogorov–Arnold Networks（KAN）

再びMLPとシベンコの定理、KANとコルモゴロフ-アルノルト表現定理の関係を強調します。

2.1 コルモゴロフ-アルノルト表現定理

コルモゴロフ-アルノルト表現定理は、有界なドメインで連続な多変数関数が有限個の一変数連続関数の合成と加法で表せるという定理です。より正確には、滑らかな $f : [0, 1]^{n} \to \mathbb{R}$ は $\phi_{q,p} : [0, 1] \to \mathbb{R}$ と $\Phi_{q} : \mathbb{R} \to \mathbb{R}$ に対して $$f\left( x_{1}, \cdots, x_{n} \right) = \sum_{q=1}^{2n+1} \Phi_{q}\left( \sum_{p=1}^{n} \phi_{q,p}\left( x_{p} \right) \right)$$ と表せるというものです。問題は、コルモゴロフ-アルノルト表現定理が単に存在性のみを述べる定理であり、これらの $\phi_{q,p}$ と $\Phi_{q}$ に何の制約もないということです。したがって、実際にこのような近似が存在するとしても、現実的にその関数を見つけられない可能性があり、機械学習の分野ではほとんど相手にされなかったとまで付け加えています。

しかし、著者らはこの定理の有用性を楽観的に見ており、数式で言う構造を正確に守ることに固執しないと述べています。また、科学や日常生活では滑らかでスパースな構造が多いため、実際にはうまく機能する可能性が高いと考えているようです。

2.2 KANのアーキテクチャ

ここからいくつかの記法が紹介されます。まず、KANの形は次のように整数の配列で表されます。 $$\left[ n_{0}, n_{1}, \cdots, n_{L} \right]$$ $n_{l}$ は第 $l$ 層のノード数であり、このKANは入力層と出力層を含めて $\left( L+1 \right)$ 個の層を持っています。第 $l$ 層の第 $i$ 番目のノードを $\left( l, i \right)$ と表し、このノードの値を $x_{l,i}$ とします。$l$ 番目の層と $(l+1)$ 番目の層の間には正確に $n_{l} n_{l+1}$ 個のリンク（活性化関数）があり、2つのニューロン $(l,i)$ と $\left( l+1, j \right)$ をつなぐリンクの活性化関数を $\phi_{l,j,i}$ と表します。ここでインデックスの範囲は当然 $l = 0, \cdots, L-1$ と前の層の $i = 1, \cdots, n_{l}$、そして後の層の $j = 1, \cdots, n_{l+1}$ で定まります。

後の層のノードの値 $x_{l+1, j}$ は $$x_{l+1, j} = \sum_{i=1}^{n_{l}} \phi_{l,j,i}\left( x_{l, i} \right)$$ と表され、これらのベクトルは $\mathbf{x}_{l} = \left( x_{l, 1}, \cdots, x_{l, n_{l}} \right)$ のようにボールド体で表記します。この過程をより理解しやすく行列形式で表すと次のようになります。 $$\begin{align*} & \mathbf{x}_{l+1} \ =& \Phi_{l}\left( \mathbf{x}_{l+1} \right) \ =& \begin{bmatrix} \phi_{l,1,1}(\cdot) & \phi_{l,1,2}(\cdot) & \cdots & \phi_{l,1,n_{l}}(\cdot) \ \phi_{l,2,1}(\cdot) & \phi_{l,2,2}(\cdot) & \cdots & \phi_{l,2,n_{l}}(\cdot) \ \vdots & \vdots & \ddots & \vdots \ \phi_{l,n_{l+1},1}(\cdot) & \phi_{l,n_{l+1},2}(\cdot) & \cdots & \phi_{l,n_{l+1},n_{l}}(\cdot) \end{bmatrix} \mathbf{x}_{l} \ =& \begin{bmatrix} \phi_{l,1,1}\left( x_{l, 1} \right) + \phi_{l,1,2}\left( x_{l, 2} \right) + \cdots + \phi_{l,1,n_{l}}\left( x_{l, n_{l}} \right) \ \phi_{l,2,1}\left( x_{l, 1} \right) + \phi_{l,2,2}\left( x_{l, 2} \right) + \cdots + \phi_{l,2,n_{l}}\left( x_{l, n_{l}} \right) \ \vdots \ \phi_{l,n_{l+1},1}\left( x_{l, 1} \right) + \phi_{l,n_{l+1},2}\left( x_{l, 2} \right) + \cdots + \phi_{l,n_{l+1},n_{l}}\left( x_{l, n_{l}} \right) \end{bmatrix} \ =& \begin{bmatrix} \sum_{i=1}^{n_{l}} \phi_{l,1,i}\left( x_{l, i} \right) \ \sum_{i=1}^{n_{l}} \phi_{l,2,i}\left( x_{l, i} \right) \ \vdots \ \sum_{i=1}^{n_{l}} \phi_{l,n_{l+1},i}\left( x_{l, i} \right) \end{bmatrix} \end{align*}$$ これは数学的に厳密な表現ではないので注意が必要です。$\Phi_{l}$ と $\mathbf{x}_{l}$ の行列積ではなく、関数 $\phi$ たちにベクトルの各成分を入れて足し合わせる過程がまるで行列の演算と似ているため、見やすく示しただけです。

この表現に従って、KANは次のように $\Phi_{l}$ たちの合成関数で表すことができます。 $$\operatorname{KAN}\left( \mathbf{x} \right) = \left( \Phi_{L-1} \circ \cdots \circ \Phi_{1} \circ \Phi_{0} \right)\left( \mathbf{x} \right)$$ これはMLPが活性化関数 $\sigma$ とアフィン変換 $W_{l}$ に対して次のように表せるのと類似しています。 $$\operatorname{MLP}\left( \mathbf{x} \right) = \left( W_{L-1} \circ \sigma \circ \cdots \circ \sigma \circ W_{1} \circ \sigma \circ W_{0} \right)\left( \mathbf{x} \right)$$

実装の詳細

$$\phi(x) = w_{b} b(x) + w_{s} \operatorname{spline}(x)$$ 活性化関数 $\phi$ は基本的に上記の形を持ちます。$w_{b}$ と $w_{s}$ はハイパーパラメータであり、必須ではなく、特に $w_{s}$ は学習過程で何の意味も持ちませんが、一応モデルを調整する要素として残しておくと述べています。活性化関数 $b(x)$ とスプライン $\operatorname{spline}(x)$ は自由に変更できる要素ですが、この研究ではとりあえず次のようにしました。 $$\begin{align*} b(x) =& \operatorname{SiLU} := x \sigma(x) = \frac{x}{1 + e^{-x}} \ \operatorname{spline} =& \sum_{i} c_{i} B_{i}(x) \end{align*}$$ ここで $\operatorname{SiLU}$ はシグモイド加速線形ユニット、$\sigma$ はロジスティック関数、$B_{i}$ はB-スプラインであり、B-スプラインの係数 $c_{i}$ たちがまさに学習の対象となる^trainable変数です。B-スプラインは与えられたドメイン外では値が $0$ になってしまうので、学習過程で同じ層にあるスプラインたちは前の層から来た値に応じてグリッドを更新していく必要性があります。これが果たして効率的かは別として、著者らが主張したように「活性化関数が学習する」という言葉自体は事実です。

2.3 KANの近似能力とスケーリング則

いよいよKANの核心となる本当の定理が登場します。

近似理論、KAT：$\mathbf{x} = \left( x_{1}, \cdots, x_{n} \right)$ とします。$(k-1)$ 回微分可能な $\phi_{l,j,i}$ に対して $f$ が次のように表されると仮定します。 $$\begin{align*} \Phi_{l} := & \left( \sum_{i=1}^{n_{l}} \phi_{l,1,i}, \cdots, \sum_{i=1}^{n_{l}} \phi_{l,n_{l+1},i} \right) \ f =& \left( \Phi_{L-1} \circ \cdots \circ \Phi_{1} \circ \Phi_{0} \right) \end{align*}$$ すると、グリッドサイズが $G$ の $k$ 次B-スプラインを $\phi_{l,j,i}^{G}$ と表すとき、すべての $0 \le m \le k$ に対して次を満たす $f$ に依存した定数 $C$ が存在します。 $$\left| f - \left( \Phi_{L-1}^{G} \circ \cdots \circ \Phi_{1}^{G} \circ \Phi_{0}^{G} \right) \right| \le C G^{-k-1+m}$$ ここでノルム $\left| \cdot \right|$ は微分作用素 $D$ に対して次のように導関数の大きさを測る方式で定義される $C^{m}$-ノルムです。 $$\left| g \right| := \max_{| \beta | \le m} \sup_{x \in [0, 1]^{n}} \left| D^{\beta} g(x) \right|$$

この定理はB-スプラインの性質によって容易に証明されるように見えますが、B-スプラインに関する背景知識がかなり要求されます。これで元のコルモゴロフ-アルノルト表現定理とは距離ができましたが、無から有を生じるわけでもなく、$\phi$ にB-スプラインを導入した意図は明確に現れています。このように誤差の上限が $G^{-k-1+m}$ で抑えられるということは、B-スプラインにより大きなコストを支払うほど関数の近似が正確になることを示唆しています。 $$\lim_{G, k \to \infty} G^{-k-1+m} = \lim_{G, k \to \infty} \frac{G^{m-1}}{G^{k}} = 0$$ 具体的には上記のようにB-スプラインで使用されるグリッドの数 $G$ を増やすほど、次数 $k$ を大きくするほど誤差の上限が減少するということです。確かにこの定理がこの論文の核心であり、新しいニューラルネットワークの理論的根拠となります。しかし、個人的に見ても非常に優れた業績だと認めますが、いずれにせよコルモゴロフ-アルノルト表現定理はモチーフだけが残り、実際の突破口はB-スプラインで見つけたことを確かめておきましょう。

一方、著者らがここでもう一度強調するのは、そのような上限が入力ベクトルのサイズ $n$ に無関係であるため、次元の呪いを克服したということですが、実際には同意し難いです。数式的におかしいというわけではありませんが、実戦的な状況では誤差の上限が大きな意味を持たないこともあり、今後KANを基盤とした変形アーキテクチャが出てくればまた話が変わるからです。実質的にこのような主張を堂々とするなら、画像データなどに対してもベンチマークを行い、次元の呪いを克服したことを立証すべきではないかと思います。

2.4 精度向上のために：グリッド拡張

当然といえば当然ですが、B-スプラインで近似された関数はグリッドをさらに追加し、つまりスプラインのドメインをより細かく分割したり、範囲外の点を追加することで、既存のモデルを維持しながら性能を改善する余地があります。重み一つの修正がモデル全体の構造にどう影響するかわからないMLPと比較すればもちろん良いことですが、実用的な側面で想像するに、ここまで細かなコントロールをする機会があるかはわかりません。実際本文でも「小さなKANはより一般化する」と述べており、果たしてグリッドを増やすことがいつも良いと見なせるかについては懐疑を示しているので、あまり深く考える必要はないでしょう。

2.5 解釈可能性のために：KANの簡素化とインタラクティブ化

先に証明された定理と同じくらい重要なコンセプトとして、KANの構造をスパースにして解釈可能性を高めるアイデアについて紹介します。

2.5.1 簡素化手法

MLPとは異なり、KANは行列積がないため、層と層の間の値がどのように伝達されるかすべて把握でき、もし影響力の小さいリンクの活性化関数がすべての値を $0$ にマッピングする定数関数 $\mathbf{0}(\cdot)$ に収束すれば、そのリンクを取り除き、人間が読める程度に明示的な公式に変える可能性も出てきます。そのために著者らは損失関数に2つの項を追加します。

1つ目はラッソ回帰のアイデアそのままに $L_{1}$ ノルム $\left| \cdot \right|_{1}$ を使用し、データサイエンスに馴染みのある人なら自然な流れとして感じられるほど直感的です。活性化関数 $\phi$ の $L_{1}$ ノルム $\left| \phi \right|_{1}$ は、$\phi$ を通過する $s$ 番目の入力値 $x^{(s)}$ に対して、$N$ 個のすべての出力値 $\phi\left( x^{(s)} \right)$ の平均として定義され、数式で表すと次のようになります。 $$\left| \phi \right|_{1} := \frac{1}{N} \sum_{s=1}^{N} \left| \phi\left( x^{(s)} \right) \right|$$ これにより、層 $\Phi$ の $L_{1}$ ノルム $\left| \Phi \right|_{1}$ も自然に次のように定義されます。 $$\left| \Phi \right|_{1} := \sum_{i=1}^{n_{\text{in}}} \sum_{j=1}^{n_{\text{out}}} \left| \phi_{j,i} \right|_{1}$$ しかし、著者らは層の $L_{1}$ ノルムだけでは十分でないことを発見し、次のように層 $\Phi$ のエントロピー $S\left( \Phi \right)$ も定義します。 $$S\left( \Phi \right) := - \sum_{i=1}^{n_{\text{in}}} \sum_{j=1}^{n_{\text{out}}} \frac{\left| \phi_{j,i} \right|_{1}}{\left| \Phi \right|_{1}} \log \frac{\left| \phi_{j,i} \right|_{1}}{\left| \Phi \right|_{1}}$$ 最後に、元の損失 $\mathscr{l}$ に重み $\mu_{1}$ と $\mu_{2}$ をかけて足した総損失 $\mathscr{L}$ を次のように定義します。 $$\mathscr{L} = \mathscr{l} + \mu_{1} \sum_{l=0}^{L-1} \left| \Phi_{l} \right|_{1} + \mu_{2} \sum_{l=0}^{L-1} S\left( \Phi_{l} \right)$$ $\phi$ の相対的な大きさで定義されるエントロピーを最小化するということは、どのような意味を持つのでしょうか。確率論に詳しい人ならば、すぐに納得できるほど巧妙な妙手であり、関数値の無秩序度を最小化するということは、それ自体で $\left| \phi \right|_{1}$ を大きくしたり $0$ に収束させる作用を追加することになります。

しかし、このように完成したモデルをシンボル化^{symbolification}する過程からは、どうしても納得できない部分が多いです。図2.4で説明しているように、まずは可視化^{visualization}から始まります。$\phi : \mathbb{R} \to \mathbb{R}$ は $x$ 軸と $y$ 軸を与えてグラフを描くことができ、$\tanh\left( 3 \left| \phi \right|_{1} \right)$ で透明度を与えてどのリンクが重要か大まかに把握しておきます。次に、ノードごとに入力スコア^{incoming score} $I_{l,i}$ と出力スコア^{outgoing score} $O_{l,i}$ を次のように定義します。 $$\begin{align*} I_{l,i} :=& \max_{k} \left| \phi_{l-1, i, k} \right| \ O_{l,i} :=& \max_{k} \left| \phi_{l+1, k, i} \right| \end{align*}$$ 2つのスコアがしきい値 $\theta = 10^{-2}$ を超えない場合、そのノードを削除します。最後に、関数の概形を見て適当な候補 $f$ を見つけた後、$\phi$ のドメインにある様々な $x$ を代入した $y = \phi(x)$ をサンプリングします。これが $y \approx c f(a x + b) + d$ を最大限よく満たす $a, b, c, d$ を見つけ、$\phi$ を $f$ で置き換えます。

🤔… もちろんこのように数式を作る過程がKANの核心ではありません。数式を見つける過程を手作業で行うからといって、研究の核心価値が下がるわけでもありません。しかし、これまでの斬新な展開に比べると、その結末がやや物足りなく見えるのも事実です。続いて著者らはシンボリック回帰^{symbolic regression}とKANを比較し、セクション2を締めくくります。シンボリック回帰は扱いにくいのに対し、KANで数式を作る過程はKANの脳を見せてユーザーが直接手術をするようなものだと述べる程度で終わります。

ここからはKANとMLPを比較したり、様々な問題に挑戦しながら実験的にKANの可能性を立証する内容が続きます。約30ページ分のすべての内容を細かく見る必要はないので、大きなポイントだけ強調してレビューを終えます。

3. KANは正確である

3.3 ファインマンデータセット

ファインマンデータセット^{Feynman dataset}は、ファインマンの著書から収集した物理方程式で作られたデータセットで、数十種類の方程式に対してKANがうまく機能することを示しています。

3.4 偏微分方程式の解法

PINNは主に偏微分方程式を解くために多く応用されており、KANとPINNが共存できることを示す例としてポアソン方程式を採用しました。PINNはニューラルネットワークのアーキテクチャがどうであれ、損失関数に方程式の情報を含めることが核心アイデアであり、数式的に見たときにKANと組み合わせられない理由はありません。

4. KANは解釈可能である

4.3 数学への応用：結び目理論

純粋数学の研究に役立つ可能性がある例として結び目理論^{knot theory}を挙げました。結び目理論では、与えられた2つの結び目が互いに同じかどうかを判別できる定理やアルゴリズムなどに多くの関心を持ち、結び目が固有に持つ不変量^invariantsを用いて結び目と結び目を区別する方法を好んで使います。自然に浮かぶ質問は「多数の結び目と多数の不変量がどのような関係を持つのか」、つまり「関数を見つけることができるのか」です。この例で著者らは「AI for Math」という表現まで使ってKANが研究に役立つことをアピールしています。

4.4 物理学への応用：アンダーソン局在化

私がアンダーソン局在化^{Anderson localization}についてよく知らないので詳しく説明するのは難しいですが、要点は結局これまでと同じです。KANが物理学の研究に役立つということです。

5. 関連研究

KANと関連する研究について論じます。論文に必要な内容ではないかもしれませんが、2024年時点で関連する分野がどのようなトレンドをたどり、KANに至るまでどのようなアイデアがあったのかについて説明します。

6. 議論

最終的な結論：KANとMLPのどちらを使うべきか？

それでは、KANとMLPのどちらを使うべきでしょうか？著者らは解釈可能性が重要であり、トレーニング時間がそれほど大きな問題でない場合、それほどスケールが大きくないAIや自然科学の問題でKANを試すことを勧めています。

余談

これまで紹介されたKANは、学習するたびにB-スプラインのドメインを更新し、新しい関数を見つけるという、従来のMLPに比べて非常に複雑な方式で実装されています。現在、MLPを基盤として広く使われているニューラルネットワークはGPUの強力な線形代数計算を積極的に使用していますが、KANは根本的にこのような物量戦に弱いと言わざるを得ません。実際、著者らは同じ数のパラメータを持つMLPモデルに比べてKANが約10倍遅く学習すると述べていますが、私が見るには10倍の速度差でさえ非常に楽観的なのではないかと思います。ディープラーニングの真価は、ヒューリスティックなニューラルネットワーク構造の変化にも飛躍的な性能改善の余地があることにあります。

ただし、KANの理論的根拠は明確であり、ディープラーニングのようなブラックボックス手法が持っていた気がかりな欠点を解決できる強力な候補であることも認めます。当然ながら、KANが登場するや否や、この速度問題を解決しようとする試みが続きました。

• efficient KAN：$\left| \phi \right|_{1}$ を計算する方式を変更しました。出力値の絶対値平均を計算する過程がかなり大きなコストを伴うため、わざわざそうするのではなく、パラメータ（B-スプラインの係数）の絶対値を使っても数式的な意味は大きく変わらないでしょう。リンクを切るという意図とは別に、ある観点ではより正確にラッソ回帰の方式を継承したと言えます。
• FastKAN：B-スプラインの代わりにガウス放射関数を使用します。ガウシアン形の関数は $k=3$ のB-スプラインと非常によく似ているため、パフォーマンスは大きく変わらないのに対して、トレーニング速度は非常に速くなることが予想されます。しかし、これは工学的に妥当な改造であるだけで、KANの理論的基盤となるB-スプラインを放棄したため、コルモゴロフ-アルノルト表現定理とはまた一段と離れたことを知っておく必要があります。