Riepilogo Tecnico: Predire Funzioni, non Caratteristiche – KAN con Apprendimento Predittivo di Embedding Congiunto nello Spazio delle Funzioni

Definizione del Problema

La segmentazione di immagini mediche affronta sfide quali il debole contrasto, i confini ambigui e le ampie variazioni di aspetto. Sebbene le architetture encoder-decoder (ad esempio, le varianti di U-Net) affrontino questi problemi tramite la fusione multi-scala, esse si affidano a pesi scalari e attivazioni fisse, lasciando le trasformazioni non lineari definite implicitamente.

Le reti Kolmogorov–Arnold (KAN) offrono una rappresentazione funzionale esplicita parametrizzando gli archi della rete come funzioni univariata apprendibili ($\phi_{j,i}$). Tuttavia, i modelli di segmentazione basati su KAN esistenti (ad esempio, U-KAN) ottimizzano queste funzioni di arco solo attraverso obiettivi definiti dopo l'aggregazione degli archi. Ciò crea due limitazioni primarie:

1. Mancanza di Target Espliciti di Pre-Aggregazione: La supervisione vincola la somma delle risposte degli archi in ingresso, permettendo a diverse configurazioni di funzioni di arco di produrre output identici ai nodi. Di conseguenza, le singole funzioni di arco mancano di un target di apprendimento esplicito prima dell'aggregazione.
2. Caratterizzazione Insufficiente del Comportamento Locale: Durante un passaggio in avanti (forward pass), una funzione viene valutata solo in un singolo punto ancora dipendente dall'input. Funzioni con risposte identiche in quel punto possono esibire comportamenti locali differenti nelle vicinanze, eppure tali variazioni rimangono implicitamente vincolate dal compito a valle.

I metodi esistenti di Apprendimento Predittivo di Embedding Congiunto (JEPA) operano su embedding di feature, token o immagini, ma non trattano le funzioni neurali apprendibili come target predittivi. Estendere JEPA allo spazio delle funzioni KAN è non banale perché una singola risposta di un arco non può caratterizzare il comportamento locale, e gli archi campionati devono rimanere esplicitamente allineati tra i rami.

Metodologia: FS-JEPA

Gli autori propongono FS-JEPA (Function-Space Joint-Embedding Predictive Learning), un framework che sposta l'apprendimento predittivo nello spazio delle funzioni pre-aggregazione delle KAN.

Framework Core

L'architettura consiste in una rete di segmentazione KAN a forma di U standard e un ramo predittivo utilizzato solo durante l'addestramento.

• Ramo Target: Un ramo a Media Mobile Esponenziale (EMA) a contesto completo genera le firme target.
• Ramo Online: Un ramo mascherato riceve una sequenza di token parzialmente mascherata e predice le firme target.
• Corrispondenza degli Archi: Entrambi i rami condividono gli stessi indici di archi campionati per garantire che le predizioni corrispondano alle stesse funzioni di arco input-output.

Componenti Chiave

1. Campionamento e Selezione degli Archi
Per evitare il costo proibitivo di calcolare le risposte per tutti gli $O(BN d_{in} d_{out})$ archi, il framework campiona casualmente $K$ coppie di archi input-output. Un operatore di selezione degli archi ($\Phi_S$) valuta solo questi specifici archi raccogliendo i loro parametri spline e le griglie dei nodi (knot grids), riducendo i costi di attivazione ausiliari a $O(BNK)$.

2. Firma della Funzione dell'Arco Multi-Raggio
Inve al posto di predire una singola risposta isolata, il metodo rappresenta ogni funzione di arco utilizzando una firma multi-raggio.

• Per un arco con input ancora $u$, viene definito un insieme di offset locali $\mathcal{D} = \{-0.10, -0.05, -0.025, 0.025, 0.05, 0.10\}$.
• La firma $s_e$ è il vettore delle valutazioni della funzione a questi offset: $s_e = [\phi_e(u+\Delta)]_{\Delta \in \mathcal{D}}$.
• Questa rappresentazione strutturata cattura le variazioni funzionali locali che una valutazione in un singolo punto non può catturare.

3. Corrispondenza Deterministica degli Archi
Per risolvere il problema del campionamento stocastico, le posizioni sorgente e target degli archi campionati $(i, j)$ sono codificate come coordinate deterministiche $(p_i, p_j)$. Queste coordinate vengono concatenate con la firma della funzione e fornite come condizioni al predittore di firma, garantendo una corrispondenza coerente senza parametri apprendibili.

4. Ottimizzazione Congiunta
Il framework ottimizza congiuntamente la perdita di segmentazione ($L_{seg}$) e la perdita predittiva nello spazio delle funzioni ($L_{sig}$).

• Perdita Predittiva: Definita come $L_{sig} = \text{SmoothL1}(\hat{s}_e, s^+_e)$, dove $\hat{s}_e$ è la predizione dal ramo online e $s^+_e$ è il target dal ramo EMA.
• Addestramento: La perdita totale è $L_{total} = L_{seg} + \lambda_{deep}L_{deep} + \lambda_{sig}\sum L^{(b)}_{sig}$.
• Inferenza: Tutte le componenti predittive (ramo online, ramo EMA, predittore di firma) vengono rimosse, lasciando invariato il percorso di segmentazione originale.

Contributi Principali

1. JEPA nello Spazio delle Funzioni: La formulazione dell'apprendimento predittivo di embedding congiunto nello spazio nativo delle funzioni degli archi delle KAN, spostando il target predittivo dagli embedding delle feature alle firme strutturate delle funzioni degli archi.
2. Rappresentazione della Firma Multi-Raggio: L'introduzione di una rappresentazione che cattura il comportamento funzionale locale delle funzioni degli archi KAN, fornendo un target predittivo efficace per l'apprendimento nello spazio delle funzioni.
3. Framework Predittivo Solo per l'Addestramento: Un metodo che apprende congiuntamente le firme delle funzioni degli archi KAN e le rappresentazioni di segmentazione da un'inizializzazione casuale mantenendo la corrispondenza degli archi, senza costi computazionali o parametri aggiuntivi durante l'inferenza.

Risultati Sperimentali

Il metodo è stato valutato su cinque benchmark di segmentazione di immagini mediche: BUSI (ultrasuoni), DDTI (ultrasuoni), TN3K (tiroide), CVC-ClinicDB (polipi) e GlaS (ghiandole).

• Prestazioni: FS-JEPA ha raggiunto il punteggio medio di Dice più elevato (83,37%) e IoU (75,04%) attraverso i cinque dataset.
• Confronto: Ha superato il più forte metodo basato su KAN concorrente (UUEKAN) di +2,25 punti percentuali in Dice.
• Ablazione:
  • Rispetto a uno scaffold corrispondente addestrato senza FS-JEPA, il metodo proposto ha migliorato il Dice medio di 0,85 punti percentuali.
  • L'ablation sulla predizione dei target ha mostrato che predire firme di funzioni di arco multi-raggio (0,7926 Dice) supera significativamente la predizione di singole risposte di arco (0,7771 Dice) o di feature di token/nodo.
  • L'analisi meccanicistica ha confermato che FS-JEPA fornisce un segnale di apprendimento pre-aggregazione esplicito che modella direttamente le rappresentazioni delle funzioni degli archi KAN, rendendole sensibili alle variazioni a livello di arco che la perdita di segmentazione standard non coglie.

Significato e Rivendicazioni

Il documento afferma che FS-JEPA dimostra l'efficacia dell'estensione dell'apprendimento predittivo oltre i convenzionali embedding di feature alle rappresentazioni strutturate delle funzioni neurali apprendibili. Predicendo firme di funzioni multi-raggio prima dell'aggregazione, il metodo consente l'apprendimento diretto delle rappresentazioni sulle funzioni degli archi KAN. Gli autori asseriscono che i miglioramenti delle prestazioni derivano dalla supervisione predittiva nello spazio delle funzioni piuttosto che da un aumento della capacità del modello, poiché il numero di parametri all'inferenza rimane identico a quello dello scaffold di base. Il framework offre un modo per vincolare esplicitamente il comportamento funzionale dei singoli archi KAN, affrontando una limitazione fondamentale nei modelli di visione basati su KAN esistenti.