Technical Summary: Predicting Functions, Not Features – KANs with Function-Space Joint-Embedding Predictive Learning

Problem Statement

Medical image segmentation faces challenges such as weak contrast, ambiguous boundaries, and large appearance variations. While encoder-decoder architectures (e.g., U-Net variants) address these via multi-scale fusion, they rely on scalar weights and fixed activations, leaving nonlinear transformations implicitly defined.

Kolmogorov–Arnold Networks (KANs) offer an explicit functional representation by parameterizing network edges as learnable univariate functions ($\phi_{j,i}$). However, existing KAN-based segmentation models (e.g., U-KAN) optimize these edge functions only through objectives defined after edge aggregation. This creates two primary limitations:

1. Lack of Explicit Pre-Aggregation Targets: The supervision constrains the sum of incoming edge responses, allowing different edge-function configurations to produce identical node outputs. Consequently, individual edge functions lack an explicit learning target before aggregation.
2. Insufficient Characterization of Local Behavior: During a forward pass, a function is evaluated only at a single input-dependent anchor. Functions with identical responses at that point may exhibit different local behaviors nearby, yet these variations remain implicitly constrained by the downstream task.

Existing Joint-Embedding Predictive Learning (JEPA) methods operate on feature, token, or image embeddings but do not treat learnable neural functions as prediction targets. Extending JEPA to the KAN function space is non-trivial because a single edge response cannot characterize local behavior, and sampled edges must remain explicitly aligned between branches.

Methodology: FS-JEPA

The authors propose Function-Space Joint-Embedding Predictive Learning (FS-JEPA), a framework that shifts predictive learning into the pre-aggregation function space of KANs.

Core Framework

The architecture consists of a standard U-shaped KAN segmentation network and a training-only predictive branch.

• Target Branch: A full-context Exponential Moving Average (EMA) branch generates target signatures.
• Online Branch: A masked branch receives a partially masked token sequence and predicts the target signatures.
• Edge Correspondence: Both branches share the same sampled edge indices to ensure predictions correspond to the same input-output edge functions.

Key Components

1. Edge Sampling and Selection
To avoid the prohibitive cost of computing responses for all $O(BN d_{in} d_{out})$ edges, the framework randomly samples $K$ input-output edge pairs. A selected-edge operator ($\Phi_S$) evaluates only these specific edges by gathering their spline parameters and knot grids, reducing auxiliary activation costs to $O(BNK)$.

2. Multi-Radius Edge Function Signature
Instead of predicting a single isolated response, the method represents each sampled edge function using a multi-radius signature.

• For an edge with anchor input $u$, a set of local offsets $\mathcal{D} = \{-0.10, -0.05, -0.025, 0.025, 0.05, 0.10\}$ is defined.
• The signature $s_e$ is the vector of function evaluations at these offsets: $s_e = [\phi_e(u+\Delta)]_{\Delta \in \mathcal{D}}$.
• This structured representation captures local functional variations that a single point evaluation cannot.

3. Deterministic Edge Correspondence
To resolve the stochastic sampling issue, the source and target node positions of a sampled edge $(i, j)$ are encoded as deterministic coordinates $(p_i, p_j)$. These coordinates are concatenated with the function signature and provided as conditions to the signature predictor, ensuring consistent correspondence without learnable parameters.

4. Joint Optimization
The framework jointly optimizes the segmentation loss ($L_{seg}$) and the function-space predictive loss ($L_{sig}$).

• Predictive Loss: Defined as $L_{sig} = \text{SmoothL1}(\hat{s}_e, s^+_e)$, where $\hat{s}_e$ is the prediction from the online branch and $s^+_e$ is the target from the EMA branch.
• Training: The total loss is $L_{total} = L_{seg} + \lambda_{deep}L_{deep} + \lambda_{sig}\sum L^{(b)}_{sig}$.
• Inference: All predictive components (online branch, EMA branch, signature predictor) are removed, leaving the original segmentation pathway unchanged.

Key Contributions

1. Function-Space JEPA: The formulation of joint-embedding predictive learning in the native pre-aggregation function space of KANs, shifting the prediction target from feature embeddings to structured edge-function signatures.
2. Multi-Radius Signature Representation: The introduction of a representation that captures local functional behavior of KAN edge functions, providing an effective prediction target for function-space learning.
3. Training-Only Predictive Framework: A method that jointly learns KAN edge-function signatures and segmentation representations from random initialization while maintaining edge correspondence, with no additional parameters or computational cost during inference.

Experimental Results

The method was evaluated on five medical image segmentation benchmarks: BUSI (ultrasound), DDTI (ultrasound), TN3K (thyroid), CVC-ClinicDB (polyps), and GlaS (glands).

• Performance: FS-JEPA achieved the highest average Dice score (83.37%) and IoU (75.04%) across the five datasets.
• Comparison: It outperformed the strongest competing KAN-based method (UUEKAN) by +2.25 percentage points in Dice.
• Ablation:
  • Compared to a matched scaffold trained without FS-JEPA, the proposed method improved average Dice by 0.85 percentage points.
  • Ablation on predictive targets showed that predicting multi-radius edge-function signatures (0.7926 Dice) significantly outperformed predicting single edge responses (0.7771 Dice) or token/node features.
  • Mechanistic analysis confirmed that FS-JEPA provides an explicit pre-aggregation learning signal that directly reshapes KAN edge-function representations, making them sensitive to edge-level variations that standard segmentation loss misses.

Significance and Claims

The paper claims that FS-JEPA demonstrates the effectiveness of extending predictive learning beyond conventional feature embeddings to the structured representations of learnable neural functions. By predicting multi-radius function signatures before aggregation, the method enables direct representation learning over KAN edge functions. The authors assert that performance improvements originate from this function-space predictive supervision rather than increased model capacity, as the inference-time parameter count remains identical to the baseline scaffold. The framework offers a way to explicitly constrain the functional behavior of individual KAN edges, addressing a fundamental limitation in existing KAN-based vision models.