Resumen Técnico: Predicción de Funciones, No de Características – KANs con Aprendizaje Predictivo de Incrustación Conjunta en el Espacio de Funciones

Planteamiento del Problema

La segmentación de imágenes médicas enfrenta desafíos como el contraste débil, los límites ambiguos y las grandes variaciones de apariencia. Aunque las arquitecturas codificador-decodificador (por ejemplo, variantes de U-Net) abordan esto mediante la fusión multiescala, dependen de pesos escalares y activaciones fijas, dejando las transformaciones no lineales definidas de forma implícita.

Las Redes Kolmogorov–Arnold (KANs) ofrecen una representación funcional explícita al parametrizar las aristas de la red como funciones univariadas aprendibles ($\phi_{j,i}$). Sin embargo, los modelos de segmentación basados en KAN existentes (por ejemplo, U-KAN) optimizan estas funciones de arista solo a través de objetivos definidos después de la agregación de las aristas. Esto crea dos limitaciones primarias:

1. Falta de Objetivos Explícitos de Pre-Agregación: La supervisión restringe la suma de las respuestas de las aristas entrantes, lo que permite que diferentes configuraciones de funciones de arista produzcan salidas de nodo idénticas. En consecuencia, las funciones de aristas individuales carecen de un objetivo de aprendizaje explícito antes de la agregación.
2. Caracterización Insuficiente del Comportamiento Local: Durante una pasada hacia adelante (forward pass), una función se evalúa solo en un único ancla dependiente de la entrada. Las funciones con respuestas idénticas en ese punto pueden exhibir comportamientos locales diferentes en las proximidades, pero estas variaciones permanecen restringidas implícitamente por la tarea de seguimiento.

Los métodos de Aprendizaje Predictivo de Incrustación Conjunta (JEPA) existentes operan sobre incrustaciones (embeddings) de características, tokens o imágenes, pero no tratan las funciones neuronales aprendibles como objetivos de predicción. Extender JEPA al espacio de funciones de KAN no es trivial porque una respuesta de arista única no puede caracterizar el comportamiento local, y las aristas muestreadas deben permanecer explícitamente alineadas entre las ramas.

Metodología: FS-JEPA

Los autores proponen el Aprendizaje Predictivo de Incrustación Conjunta en el Espacio de Funciones (FS-JEPA), un marco que desplaza el aprendizaje predictivo hacia el espacio de funciones de pre-agregación de las KANs.

Marco Central

La arquitectura consiste en una red de segmentación KAN con forma de U estándar y una rama predictiva exclusiva para el entrenamiento.

• Rama de Objetivo (Target Branch): Una rama de Media Móvil Exponencial (EMA) de contexto completo genera las firmas de objetivo.
• Rama en Línea (Online Branch): Una rama enmascarada recibe una secuencia de tokens parcialmente enmascarada y predice las firmas de objetivo.
• Correspondencia de Aristas: Ambas ramas comparten los mismos índices de aristas muestreadas para asegurar que las predicciones correspondan a las mismas funciones de arista de entrada-salida.

Componentes Clave

1. Muestreo y Selección de Aristas
Para evitar el costo prohibitivo de computar las respuestas para todas las $O(BN d_{in} d_{out})$ aristas, el marco muestrea aleatoriamente $K$ pares de aristas de entrada-salida. Un operador de selección de aristas ($\Phi_S$) evalúa solo estas aristas específicas recopilando sus parámetros de spline y rejillas de nudos (knot grids), reduciendo los costos de activación auxiliares a $O(BNK)$.

2. Firma de Función de Arista de Radio Múltiple
En lugar de predecir una respuesta aislada única, el método representa cada función de arista muestreada utilizando una firma de radio múltiple.

• Para una arista con una entrada ancla $u$, se define un conjunto de desplazamientos locales $\mathcal{D} = \{-0.10, -0.05, -0.025, 0.025, 0.05, 0.10\}$.
• La firma $s_e$ es el vector de evaluaciones de la función en estos desplazamientos: $s_e = [\phi_e(u+\Delta)]_{\Delta \in \mathcal{D}}$.
• Esta representación estructurada captura las variaciones funcionales locales que una evaluación de un solo punto no puede capturar.

3. Correspondencia de Aristas Determinista
Para resolver el problema del muestreo estocástico, las posiciones de origen y destino de una arista muestreada $(i, j)$ se codifican como coordenadas deterministas $(p_i, p_j)$. Estas coordenadas se concatenan con la firma de la función y se proporcionan como condiciones al predictor de firmas, asegurando una correspondencia consistente sin parámetros aprendibles.

4. Optimización Conjunta
El marco optimiza conjuntamente la pérdida de segmentación ($L_{seg}$) y la pérdida predictiva de espacio de funciones ($L_{sig}$).

• Pérdida Predictiva: Definida como $L_{sig} = \text{SmoothL1}(\hat{s}_e, s^+_e)$, donde $\hat{s}_e$ es la predicción de la rama en línea y $s^+_e$ es el objetivo de la rama EMA.
• Entrenamiento: La pérdida total es $L_{total} = L_{seg} + \lambda_{deep}L_{deep} + \lambda_{sig}\sum L^{(b)}_{sig}$.
• Inferencia: Todos los componentes predictivos (rama en línea, rama EMA, predictor de firmas) se eliminan, dejando la vía de segmentación original sin cambios.

Contribuciones Clave

1. JEPA en el Espacio de Funciones: La formulación del aprendizaje predictivo de incrustación conjunta en el espacio de funciones de pre-agregación nativo de las KANs, desplazando el objetivo de predicción de las incrustaciones de características a firmas de funciones de arista estructuradas.
2. Representación de Firma de Radio Múltiple: La introducción de una representación que captura el comportamiento funcional local de las funciones de arista de KAN, proporcionando un objetivo de predicción efectivo para el aprendizaje en el espacio de funciones.
3. Marco Predictivo Solo para Entrenamiento: Un método que aprende conjuntamente las firmas de las funciones de arista de KAN y las representaciones de segmentación desde una inicialización aleatoria manteniendo la correspondencia de las aristas, sin costo computacional ni parámetros adicionales durante la inferencia.

Resultados Experimentales

El método fue evaluado en cinco benchmarks de segmentación de imágenes médicas: BUSI (ultrasonido), DDTI (ultrasonido), TN3K (tiroides), CVC-ClinicDB (pólipos) y GlaS (glándulas).

• Desempeño: FS-JEPA alcanzó el coeficiente Dice promedio más alto (83.37%) e IoU (75.04%) en los cinco conjuntos de datos.
• Comparación: Superó al método basado en KAN más fuerte (UUEKAN) por +2.25 puntos porcentuales en Dice.
• Ablación:
  • Comparado con un andamiaje (scaffold) emparejado entrenado sin FS-JEPA, el método propuesto mejoró el Dice promedio en 0.85 puntos porcentuales.
  • La ablación sobre los objetivos predictivos mostró que predecir firmas de funciones de arista de radio múltiple (0.7926 Dice) superó significativamente a la predicción de respuestas de arista única (0.7771 Dice) o de características de token/nodo.
  • El análisis mecanístico confirmó que FS-JEPA proporciona una señal de aprendizaje de pre-agregación explícita que moldea directamente las representaciones de las funciones de arista de KAN, haciéndolas sensibles a las variaciones a nivel de arista que la pérdida de segmentación estándar omite.

Significado y Reivindicaciones

El artículo afirma que FS-JEPA demuestra la efectividad de extender el aprendizaje predictivo más allá de las incrustaciones de características convencionales hacia las representaciones estructuradas de las funciones neuronales aprendibles. Al predecir firmas de funciones de radio múltiple antes de la agregación, el método permite el aprendizaje de representación directa sobre las funciones de arista de KAN. Los autores sostienen que las mejoras en el desempeño se originan de esta supervisión predictiva en el espacio de funciones y no de un aumento en la capacidad del modelo, ya que el recuento de parámetros en el tiempo de inferencia es idéntico al del andamiaje base. El marco ofrece una forma de restringir explícitamente el comportamiento funcional de las aristas de KAN individuales, abordando una limitación fundamental en los modelos de visión basados en KAN existentes.