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

arXiv: https://arxiv.org/abs/2608.12050 · HTML: https://mlsanity.com/paper/2608.12050/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 7 of its 18 references that are in the map.
- First posted on arXiv 12 Aug 2026.

## Builds on (most cited first, 7 of 7)

- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Deeply-Supervised Nets](https://mlsanity.com/paper/1409.5185/) (2014)
- [Boundary loss for highly unbalanced segmentation](https://mlsanity.com/paper/1812.07032/) (2018)
- [UNeXt: MLP-based Rapid Medical Image Segmentation Network](https://mlsanity.com/paper/2203.04967/) (2022)
- [U-KAN Makes Strong Backbone for Medical Image Segmentation and Generation](https://mlsanity.com/paper/2406.02918/) (2024)
- [MC-JEPA: A Joint-Embedding Predictive Architecture for Self-Supervised Learning of Motion and Content Features](https://mlsanity.com/paper/2307.12698/) (2023)

## Abstract (arXiv)

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.