Kolmogorov-Arnoldネットワークの開発
(github.com/KindXiaoming)- pykanは論文「KAN: Kolmogorov-Arnold Networks」と「KAN 2.0: Kolmogorov-Arnold Networks Meet Science」のためのGitHubリポジトリで、KANの学習・チュートリアル・ドキュメント・例を提供する
- KANはMLPの代替として提示されており、MLPが普遍近似定理に基づくのと同様に、KANはKolmogorov-Arnold表現定理に基づく
- 構造的には、MLPはノードに活性化関数を持ち、KANはエッジに活性化関数を持つ。この違いがモデルの精度と解釈可能性を改善し得ると説明されている
- 記号計算ブランチを使わない機械学習ユーザーは、学習前に
model.speed()を呼び出す必要がある。そうしないと、並列化されていない symbolic branch のため非常に遅くなる可能性がある - この実装は科学関連の小規模問題を主な対象としており、機械学習タスクにそのまま差し込んで使えるプラグインとは見なしにくく、ハイパーパラメータ調整と適用ごとの手法が必要になる
pykanとKANの概要
- pykanは “KAN: Kolmogorov-Arnold Networks” と “KAN 2.0: Kolmogorov-Arnold Networks Meet Science” のためのリポジトリ
- クイックスタートは hellokan、追加の例は tutorials、ドキュメントは公式ドキュメントで提供されている
- KANs はMulti-Layer Perceptrons（MLPs）の有望な代替として紹介されている
- MLPは universal approximation theorem に基づく
- KANは Kolmogorov-Arnold representation theorem に基づく
- KANとMLPは双対的な構造として説明される
- KANはエッジに活性化関数を持つ
- MLPはノードに活性化関数を持つ
- この構造の変化により、KANの精度と解釈可能性がMLPより良くなり得ると説明されている
インストールと実行環境
- pykanはPyPIまたはGitHubからインストールできる
- 前提条件は Python 3.9.7以上 と
pip - 開発者向けインストール方法:
git clone https://github.com/KindXiaoming/pykan.gitcd pykanpip install -e .
- GitHubからのインストール:
pip install git+https://github.com/KindXiaoming/pykan.git
- PyPIからのインストール:
pip install pykan
- 主な要求パッケージには
matplotlib、numpy、scikit_learn、sympy、torch、tqdm、pandas、seaborn、pyyamlなどが含まれる - Condaユーザーは
python=3.9.7環境を作成したうえで、GitHubまたはPyPIの方法でインストールできる
性能モードと計算要件
- 機械学習ユーザーが自分で学習ループを書き、
model.fit()を使わず、symbolic branch も使わない場合は、学習前にmodel.speed()を呼び出すことが重要 model.speed()を呼び出さないとsymbolic branchが有効のままで、記号計算が並列化されないため非常に遅くなる可能性がある- tutorials の例は通常、単一CPUで10分以内に実行できる
- 論文に含まれるすべての例は、単一CPUで1日以内に実行できる
- PDE向けのKAN学習は最もコストが高く、単一CPUで数時間から数日かかることがある
- モデル学習にCPUを使用した理由は、MLPとKANのPareto Frontierを得るために、数千個の小さなモデルに対してパラメータスイープを実行したため
- 問題規模が大きいタスクでは GPUの使用 が推奨される
KANのハイパーパラメータ調整
- MLPや他のネットワークで得た直感が、KANにそのまま適用できるとは限らない
- 基本的な助言は、単純な設定から始めること
- 小さなKAN shape
- 小さなgrid size
- 小さなデータ
- 正則化なし、
lamb=0
- 例えば入力5個・出力1個のタスクでは、
KAN(width=[5,1,1], grid=3, k=3)のような非常に単純な設定から試せる - うまく動かない場合はまずwidthを増やし、それでもだめならdepthを増やす流れが推奨される
- 性能が許容できる水準に達したら、より正確、またはより解釈可能なKANへと精緻化できる
- 精度を重視するなら grid extension 手法を試せるが、過学習には注意が必要
- 解釈可能性を重視するなら、
model.train(lamb=0.01)のような方法でネットワークを疎にできるlambは段階的に増やしていく方法が推奨される- 学習後のplotで明らかに不要なニューロンが見えたら、
pruned_model = model.prune()で枝刈りしたモデルを得られる - その後、精度または疎性のために追加学習したり、symbolic regressionを実行したりできる
- 精度、解釈可能性、パラメータ効率は常に相反する関係にあるわけではなく、場合によっては正の相関があることも、tradeoffがあることもある
- train/test lossの間に大きな差がある場合は、データを増やすかモデルを小さくする方向を検討すべき
gridはwidthより重要なので、まずgridを減らし、その次にwidthを減らす順序が提案されている
- 単純なモデルから始め、まずunderfitting状態を確認し、段階的に拡張して適切な領域へ移行する方法が推奨される
適用範囲と限界
- コードは数学・物理の例のような小規模な科学問題を念頭に設計されている
- 効率性と再利用性を大きく考慮していないため、その側面への批判は受け入れるとしている
- 科学的発見と科学計算に関心のあるユーザーを本来の対象としており、リポジトリも主にこの目的を維持する予定
- 効率改善の実装として efficientkan と FourierKAN が挙げられている
- 機械学習中心のユーザーにとって、KANはまだ単純な out-of-the-boxプラグイン ではない
- ハイパーパラメータ調整が必要
- 適用ごとの特殊な手法が追加される場合がある
- GraphKAN は、KANをlatent spaceで使うほうがよいと提案しており、入力後と出力前にembedding/unembedding linear layerが必要だと言及されている
- KANRL は、強化学習で学習の安定性を高めるため、一部の学習可能パラメータを固定するほうがよいと提案している
- KANが次世代LLMになるかについては、確かな直感はないとしている
- KANは高い精度と解釈可能性を重視する応用のために設計されている
- LLMにおける解釈可能性と科学における解釈可能性は大きく異なる可能性がある
- 論文の結論をLLMや一般的な機械学習タスクへ直接移すのは難しいと見ている
- KANとMLPは互いを置き換えられるものではなく、それぞれ特定の設定では長所と限界を持つ
1件のコメント
Hacker Newsの意見
論文をざっと読んで、さらに単純化してみたくなり、PyTorchレイヤーを作ってみた: https://github.com/GistNoesis/FourierKAN/
核心は本当に数行だけ。論文ではより小規模なものを念頭に置いたようなコードで、1次元関数を表現するためにスプライン補間を使い、その結果を合算している
代わりに各座標の関数を補間するのにフーリエ係数を使う別の表現を選んだ。Kolmogorov-Arnoldネットワークの表現力をつかむ助けになりそう。スプライン版より収束は容易かもしれないが、演算数はスプラインのほうが少ない
もちろん、私のコードが動かないからといって、論文側が動かないという意味ではない。望むなら実験して論文にしてもよい
フーリエ係数も並列化しやすく書きやすいので試してみたが、学習の挙動は放射基底関数のほうが良かった
efficientKANは、すべてのエッジ活性化に共通の基底関数を先に計算し、出力はその基底の線形結合として計算する
基底関数がフーリエなら、KANレイヤーは固定重みの線形レイヤー + Sin活性化 + 学習可能な重みの線形レイヤーと見なせ、これはSirenの特殊な形である
KANとMLPの間のつながりを示す例かもしれない
「私のコードが動かないからといって、論文側が動かないという意味ではない」とは言うが、実際に動作するのかが気になる
著者らのJupyterノートブックを少し触ってみたが、個人的にはExample_3_classfication.ipynb(https://github.com/KindXiaoming/pykan/blob/master/tutorials/Example_3_classfication.ipynb)が最も有用だった
著者らが選んだパラメータでは説明どおりに動くが、チュートリアル後半の分類設定でネットワーク形状を(2, 2)から(2, 2, 2)に変えると汎化に失敗する
学習損失は1e-9まで下がるが、テスト損失は3e-1付近にとどまり、より大きなネットワークにしても役に立たなかった
パラメータとデータの複雑度がはるかに大きい例が必要で、実際に学習可能なのかも見たい。MNISTが良い出発点になりそう
更新: 学習データセットのサイズを100倍に増やしたところ過学習は減ったが、今度は学習損失を1e-2未満に下げられなくなっている。実験は継続中で、GPUアクセラレーションが切実に必要。今はCPU速度が進行を制限している
変更点は3つ。学習セットを1,000個から100kサンプルに増やして過学習を解決し、データ生成時のノイズを0.1から0.07に少し下げてクラスが重ならないようにした
最も重要でKANに特化した部分は、grid=5で30ステップ学習した後、以前のモデルから初期化してgrid=10で30ステップ、さらにgrid=20で30ステップ学習したこと。これはKANで一般的なやり方で、Example_1_function_fitting.ipynb(https://github.com/KindXiaoming/pykan/blob/master/tutorials/Example_1_function_fitting.ipynb)に載っている
全体的な印象は、動作はすること、参照実装は非常に遅いためGPU実装が不可欠であること、MLP + ReLUより非線形性が強い感じで学習安定性が低いこと
うまくスケールする保証はまだなく、このアプローチでMNISTを解けるのかは必ず見たい。引き続き注視するつもり
Transformer、LSTM、ADAMのような大きなアイデアを、y=sin(x)曲線に対する50個の数字からなる学習データだけでテストしていたら、そうしたアイデアを誤って捨てていた可能性が高い
GPUで動かせば常に速いわけではなく、特に分岐が多い場合は驚くことではない
残念ながら関連テンソルがすべて正しいデバイスに載っていなかったため、KAN.pyとKANLayer.pyを修正する必要があった。一部のフォーマットを見ると、以前はdevice引数があったらしい痕跡も見える
古典的な統計にはKolmogorov-Arnoldに触発されたGAM(https://en.wikipedia.org/wiki/Generalized_additive_model)というモデルがあり、HastieとTibshiraniがGLM(https://en.wikipedia.org/wiki/Generalized_linear_model)の拡張として開発した
GLMはロジスティック回帰、線形回帰、その他いくつかの人気のある回帰モデルを一般化する
学習された基底関数を使うニューラルネットワークGAMもすでに提案されているので、今回の新しい論文で先行研究が言及されていない点は少し意外。以前の応用は解釈可能性により重点を置いていた
ニューラルネットワークの成功はスケーラビリティと密接に結びついている。アルゴリズム自体がより多くのレイヤーへ拡張できるだけでなく、ハードウェアとも相性がよくなければならない
ニューラルネットワークは大半が行列乗算で構成され、GPU には行列乗算専用のアクセラレーションがある。AlexNet が大きな影響を与えた理由も、ニューラルネットワークを GPU に載せてスケールさせ、高速化できることを示したからだ
このアルゴリズムがどれほどよくスケールするかは、論文だけでは明らかではない。アルゴリズム面でレイヤーが増えてもよく学習できるのか、ハードウェアアクセラレーションをうまく活用できるのか、どちらも不確かだ
特に、重みごとに活性化関数が付く構造が高速な行列乗算アクセラレーションを活用できるのかはよく分からない
小規模ではうまく動き、良い特性を持つ興味深いアイデアだが、ImageNet や LLM のような用途に適したアーキテクチャかどうかはまだ分からない
高速な行列乗算アクセラレーションは、当初はグラフィックスのような特定の問題に有用だったため、ハードウェアで実装されたものだ
重みごとの活性化関数が本当にうまく効くなら、人々はすぐにそれをハードウェア上で実行する方法を見つけるだろう
よくある「Transformer をあれこれ変えて、あれこれのベンチマークで少し良くなった」という類ではない新しい AI 研究を見るのは新鮮だ
そうした漸進的な改善論文も重要だが、皆少し疲れ始めているし、逸話的な証拠や最近の研究を見ると、Transformer 固有の根本的限界に近づいているようにも思えるので、新しい代替案が必要かもしれない（https://news.ycombinator.com/item?id=40179232）
今回の仕事で最も良い点は、二者択一ではないことだ。提案されている学習可能なスプライン補間活性化関数は、既存の深層ニューラルネットワークにも入れて表現力を高められる
あとは実際によりうまく動くかをテストすればよい
PR が強すぎる力になっており、以前から存在はしていたが、今はさらに影響力が増したように思う
こうした記事に賛成票を投じたり、自分が査読する立場なら最高性能の更新だけに集中しないようにすることで対抗できる。その基準はすでにゲーム化されており、私たちを間違った方向へ導いているのは明らかだ
今の基準で見れば、その本に書かれていた内容は笑ってしまうほど基本的だったが、動機として示されていたのはKolmogorov 表現定理だった。適切な活性化関数を持つ適当な3層ネットワークは、任意の連続な m-to-n 関数を表現できるという内容だ
そのためか、当時の研究の大半は3層ネットワークに集中し、Sigmoid 活性化が主流で、勾配消失が主要な問題だった
1990年代の AI 冬の時代の後、AlexNet がニューラルネットワーク研究を復活させるまで20年かかった
プレプリントを見ると、入力次元100個を「高い」と見なしており、扱っている問題の大半は入力次元が5個以下だ
私が見てきた物理学に着想を得た機械学習の設定では典型的な姿だ
次のステップは MNIST で示すことだろうが、MNIST の784次元も現代の基準では非常に小さい
しかし、そうした問題の大半では決定木が依然としてニューラルネットワークと競争力があるか、むしろよりうまくいく
興味深い。Kolmogorov ニューラルネットワークは不連続関数を表現できるが（https://arxiv.org/abs/2311.00049）、実際の適用可能性がどの程度あるのか気になっていた
このリポジトリは、それでもある程度有用であることを示しているように見える
論文にも「不連続な有界関数および非有界関数の場合、g を実用的に構成する方法はまだ知られていない」とある
OP の arXiv リンク（https://arxiv.org/abs/2404.19756）を見ると、彼らはスプラインを使っている
依然として興味深く、潜在的に有用ではあるが、追加の発見がなければ不連続関数には有用ではない。もし私が間違っているならリンクを教えてほしい。非常に関心のあるテーマだ
早とちりかもしれないが、B-スプラインの線形結合は、別のより高次の B-スプラインではないのか？
結局、関数に高次の B-スプラインをフィットしているだけなのではないかと思う
興味深いことに、このアプローチと MLP の基盤は、約66年前のほぼ同じ時期に発明または発見されたものです
1957年: https://en.wikipedia.org/wiki/Kolmogorov%E2%80%93Arnold_representation_theorem
1958年: https://en.wikipedia.org/wiki/Multilayer_perceptron
もう一つの利点は、このアプローチには局所活性化関数の係数という1種類のパラメータしかない一方で、MLP には重み、バイアス、グローバルに同一の活性化関数という3種類のパラメータがある点です
みんな Transformer の話をしていますが、私はこのアプローチを使った拡散モデルを見てみたいです
線形和の重みとスプラインの係数の間にも、大きな違いはなさそうです
自己アテンションとクロスアテンションを入れた U-Net、Vision Transformer、Diffusion Transformer などがあります
簡潔さ、つまりオッカムの剃刀の観点から主張することはできそうですが、そういう考えなのか、それとも別の理由があるのか気になります
誰かがスプラインを決定木に無理やり押し込んだような感じです
PDF の2ページ目には「KAN のノードは非線形性を適用せず、入力信号を単純に合計する」とあります