Search papers, labs, and topics across Lattice.
This paper introduces Function-Space Joint-Embedding Predictive Learning (FS-JEPA) to enhance Kolmogorov-Arnold Networks (KANs) for medical image segmentation by allowing edge functions to be optimized with explicit pre-aggregation learning targets. By predicting structured signatures of KAN edge functions rather than isolated responses, FS-JEPA captures local functional variations more effectively, leading to improved segmentation accuracy. Experimental results demonstrate that FS-JEPA outperforms the leading KAN-based method by 2.25 percentage points in average Dice scores across five medical imaging benchmarks.
FS-JEPA redefines how we optimize KANs by predicting structured edge function signatures, resulting in a significant boost in medical image segmentation performance.
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.
