Predicting Functions, Not Features: KANs with Function-Space Joint-Embedding Predictive Learning for Medical Image Segmentation

2026-08-12

AI Summary

This work addresses a key limitation in existing Kolmogorov–Arnold Network (KAN)-based medical image segmentation methods, which optimize only after aggregating edge functions and lack explicit pre-aggregation learning objectives for individual edge functions. To overcome this, the study introduces structured prediction learning into KAN’s function space for the first time, proposing a Function-Space Joint Embedding Predictive Architecture (FS-JEPA). The framework captures local functional variations via multi-radius function signatures and employs a masked online branch to predict edge function signatures generated by a full-context EMA target branch, jointly optimizing this predictive objective with the segmentation loss. Evaluated across five medical image segmentation benchmarks, the method achieves state-of-the-art average Dice scores, outperforming the strongest KAN baseline by 2.25 percentage points.

Abstract

Kolmogorov--Arnold Networks (KANs) introduce explicit functional representations by parameterizing each network edge as a learnable univariate function. However, existing KAN-based segmentation models optimize edge functions only through objectives defined after edge aggregation, leaving individual functions without an explicit pre-aggregation learning target. To address this limitation, we propose Function-Space Joint-Embedding Predictive Learning (FS-JEPA) for medical image segmentation. Our FS-JEPA framework moves predictive learning into the pre-aggregation function space of KANs. A masked online branch predicts structured signatures of sampled KAN edge functions generated by a full-context exponential moving average target branch, while shared edge indices preserve correspondence between predictions and targets. Rather than predicting an isolated edge response, we represent each sampled edge function using a multi-radius signature composed of function evaluations around its input anchor. This structured representation captures local functional variations that cannot be characterized by a single response and provides a more informative predictive target. The function-space objective is jointly optimized with the segmentation loss during training, while the predictive branch is removed at inference. Experiments on five medical image segmentation benchmarks show that our FS-JEPA achieves the best average Dice and outperforms the strongest competing KAN-based method by +2.25 percentage points.

Problem

• Kolmogorov–Arnold Networks
• medical image segmentation
• function-space learning
• predictive learning
• edge function optimization

Innovation

• Kolmogorov–Arnold Networks
• function-space predictive learning
• joint-embedding
• medical image segmentation
• structured function signature