[論文レビュー] U-KAN Makes Strong Backbone for Medical Image Segmentation and Generation

Chenxin Li, Xinyu Liu+6 more|arXiv (Cornell University)|Jun 5, 2024

Brain Tumor Detection and Classification被引用数 43

ひとこと要約

U-KAN は Kolmogorov-Arnold Networks (KAN) を U-Net バックボーンに組み込み、医用画像分割と拡散ベースの画像生成を強化し、精度と効率を向上させ、解釈性を改善する。

ABSTRACT

U-Net has become a cornerstone in various visual applications such as image segmentation and diffusion probability models. While numerous innovative designs and improvements have been introduced by incorporating transformers or MLPs, the networks are still limited to linearly modeling patterns as well as the deficient interpretability. To address these challenges, our intuition is inspired by the impressive results of the Kolmogorov-Arnold Networks (KANs) in terms of accuracy and interpretability, which reshape the neural network learning via the stack of non-linear learnable activation functions derived from the Kolmogorov-Anold representation theorem. Specifically, in this paper, we explore the untapped potential of KANs in improving backbones for vision tasks. We investigate, modify and re-design the established U-Net pipeline by integrating the dedicated KAN layers on the tokenized intermediate representation, termed U-KAN. Rigorous medical image segmentation benchmarks verify the superiority of U-KAN by higher accuracy even with less computation cost. We further delved into the potential of U-KAN as an alternative U-Net noise predictor in diffusion models, demonstrating its applicability in generating task-oriented model architectures. These endeavours unveil valuable insights and sheds light on the prospect that with U-KAN, you can make strong backbone for medical image segmentation and generation. Project page:\url{https://yes-u-kan.github.io/}.

研究の動機と目的

• より良い非線形モデリングと解釈性を備えた医用画像分割のための改良バックボーンを動機づける。
• KAN を U-Net パイプラインに統合するためのトークン化された KAN ブロックを導入する。
• 多様な医用データセットで優れた分割精度を実証する。
• 拡張して拡散ベースの画像生成を実現し、その生成能力を評価する。

提案手法

• 2 段階のアーキテクチャを提案する：Convolution Phrase に続いて Tokenized KAN Phrase、次に skip 連結を持つ U 形デコーダー。
• 畳み込み特徴のトークン化を用いてパッチを作成し、残差接続と層正規化を備えた積み重ねられたKAN層で処理する。
• 中間特徴を D 次元埋め込みに射影し、K 個の層を持つ KAN (Phi_i) で処理し、続いて DwConv と BN を適用するトークン化された KAN ブロックを実装する。
• 時刻埋め込みを注入して Diffusion U-KAN に拡張し、DDPM フレームワークでノイズ epsilon_t を予測するよう学習する。時刻条件付き KAN と MSE 損失 (L_diff) を使用。
• 3 つの医用データセットで binary cross-entropy と dice loss の組み合わせで分割モデルを訓練し、IoU と F1 スコアで評価します。効率 (Gflops, Params) を報告。
• FID と IS 指標を用いて拡散バックボーンを標準的な U-Net 変種と比較評価する。

実験結果

リサーチクエスチョン

• RQ1 Can integrating Kolmogorov-Arnold Networks into a U-Net backbone improve segmentation accuracy on medical images without a large computational burden?
• RQ2 Does a tokenized KAN block provide interpretability and performance benefits over traditional CNN/MLP/Transformer backbones in medical imaging?
• RQ3 Can U-KAN serve effectively as a diffusion-model backbone for task-oriented image generation beyond segmentation?

主な発見

• Seg-U-KAN は複数のベースラインより高い分割精度を達成する（例：BUSI、GlaS、CVC データセットでの IoU/F1）。
• Seg. U-KAN の平均分割指標は IoU 78.69 および F1 87.22 に達し、データセット全体で競争力のある効率性（Gflops 14.02、Params 6.35）を示す。
• アブレーションにより、3 層の KAN がテストされた構成の中で最良の分割性能を示す（IoU 66.65、F1 79.75）。
• KAN 層をMLPで置換すると性能が低下し、このタスクにおけるKANブロックの有効性を強調している。
• Diffusion U-KAN は標準の Diffusion U-Net 変種より生成指標を改善し、BUSI、GlaS、CVC データセットでの FID/IS が改善される（例：KANBlock を組み込んだ Diffusion U-KAN は有利な FID/IS を示す）
• 3-KAN 層構成と時刻埋め込みの拡張が、分割と拡散生成の両方の性能を向上させる。
• モデルのスケーリング研究は、より大きな U-KAN バリアントが IoU/F1 のさらなる向上をもたらす一方で Gflops の増加というコストを伴うことを示す。