Resumen Técnico: KANLib – Una implementación modular, extensible y rápida de las Redes Kolmogorov-Arnold

Planteamiento del Problema

Las Redes Kolmogorov-Arnold (KANs) han surgido como una alternativa prometedora a los Perceptrones Multicapa (MLPs) tradicionales al reemplazar los pesos lineales fijos con funciones univariantes aprendibles. Si bien las KANs ofrecen ventajas teóricas en interpretabilidad y expresividad, su adopción práctica se ve obstaculizada por tres desafíos principales:

1. Altos Costos Computacionales: Las implementaciones existentes suelen sufrir de evaluaciones de funciones de base ineficientes y operaciones tensoriales redundantes.
2. Soporte de Características Inconsistente: Diferentes marcos de trabajo (por ejemplo, PyKAN, EfficientKAN, FastKAN) ofrecen conjuntos de características dispares, lo que dificulta la realización de comparaciones y pruebas de rendimiento rigurosas.
3. Falta de Modularidad: Las implementaciones actuales suelen estar estrechamente acopladas a diseños arquitectónicos o funciones de base específicos, lo que limita la capacidad de explorar nuevas variaciones o arquitecturas híbridas.

Metodología

El artículo presenta KANLib, un marco de software unificado diseñado para sintetizar las innovaciones centrales de tres implementaciones prominentes de KAN: PyKAN (la implementación oficial de Liu et al.), EfficientKAN y FastKAN.

Diseño Arquitectónico

KANLib está construido sobre una arquitectura modular que soporta dos tipos principales de funciones de base:

• B-splines: Siguiendo la formulación original de KAN, donde las funciones aprendibles son combinaciones ponderadas de funciones de base B-spline y una rama residual (típicamente SiLU).
• Funciones de Base Radial Gaussiana (GRBFs): Adoptadas de FastKAN para reducir la complejidad computacional evitando los cálculos recursivos de B-spline.

Características Técnicas Clave

Para abordar la fragmentación de las herramientas existentes, KANLib unifica las siguientes capacidades dentro de un único entorno basado en PyTorch:

• Capas Lineales Unificadas: Una interfaz consistente para capas tanto de B-spline como de GRBF, permitiendo comparaciones arquitectónicas directas.
• Control de Grano Fino: Los usuarios pueden habilitar o deshabilitar selectivamente componentes definidos en la ecuación estándar de KAN, como la rama residual ($w_b b(x)$) o el peso de la spline ($w_s$).
• Reescalado de Rejilla Adaptativa: Extiende el mecanismo de rejilla adaptativa encontrado en PyKAN/EfficientKAN tanto a las capas de B-spline como a las de GRBF (este último restringido a rejillas equidistantes).
• Extensión de Rejilla: Soporta el aumento de la resolución de la spline durante el entrenamiento sin perder la información aprendida, una característica que anteriormente se omitía en EfficientKAN y FastKAN.
• Estrategias de Normalización: Integra la Normalización de Capa (Layer Normalization) para mapear las entradas al rango activo de las funciones de base, compensando la falta de reescalado adaptativo en ciertas configuraciones.

Contribuciones Clave

1. Unificación de Marcos de Trabajo: KANLib consolida los conjuntos de características de PyKAN, EfficientKAN y FastKAN en un único código base extensible, facilitando la comparación rigurosa entre diferentes funciones de base y estrategias de optimización.
2. Eficiencia Computacional: Al integrar optimizaciones de EfficientKAN (por ejemplo, evitar la expansión del tensor de entrada) y FastKAN (por ejemplo, el cálculo directo de GRBF), KANLib logra velocidades de inferencia competitivas manteniendo la riqueza de características.
3. Flexibilidad Arquitectónica: El marco permite la exploración sistemática de variaciones de KAN (por ejemplo, eliminar ramas residuales o pesos de spline) con cambios mínimos en el código, apoyando la investigación más allá de las formulaciones estándar.
4. Reproducibilidad: El marco está diseñado para reproducir el comportamiento predictivo de las implementaciones de referencia establecidas, sirviendo como una base de referencia fiable para investigaciones futuras.

Resultados Experimentales

Los autores evaluaron KANLib en el conjunto de datos California Housing (una tarea de regresión con 20,640 muestras y 8 características de entrada), comparándolo con PyKAN, EfficientKAN y FastKAN.

Rendimiento Predictivo

• Variantes de B-spline: KANLib (B-spline) logró el mejor rendimiento entre las implementaciones de B-spline, con un RMSE de 0.5376 ± 0.0044 y un $R^2$ de 0.7852 ± 0.0035. Esto superó a PyKAN (RMSE 0.5434) y EfficientKAN (RMSE 0.5463), demostrando que KANLib reproduce con éxito y mejora ligeramente los comportamientos basados en splines existentes.
• Variantes de GRBF: KANLib (GRBF) logró un RMSE de 0.5471 ± 0.0078, igualando de cerca a FastKAN (RMSE 0.5506), confirmando la capacidad del marco para replicar las KAN basadas en RBF gaussianas.

Eficiencia Computacional

• Tiempo de Inferencia:
  • PyKAN fue el más lento (212.19 ms).
  • EfficientKAN y KANLib (B-spline) redujeron el tiempo de inferencia aproximadamente un 32.7% en comparación con PyKAN (~142 ms).
  • KANLib (GRBF) logró 80.58 ms, una mejora del 62.0% sobre PyKAN.
  • FastKAN siguió siendo el más rápido (41.02 ms) debido a su implementación altamente especializada y menos flexible.
• Recuento de Parámetros: Los modelos de KANLib generalmente contenían 4,050 parámetros (incluyendo coeficientes de spline, pesos residuales y pesos de spline), mientras que FastKAN tenía menos (3,811) debido a la omisión del mecanismo de ponderación de spline.

Exploración Arquitectónica

Los experimentos que variaron la estructura de la capa (por ejemplo, "Sin Residual", "Sin Peso de Spline", "Plana") mostraron que:

• Eliminar la rama residual o los pesos de spline redujo el recuento de parámetros (hasta 3,510) y los tiempos de inferencia (hasta un ~16.8% más rápido para los modelos GRBF).
• El rendimiento predictivo se mantuvo competitivo en todas estas variaciones, con desviaciones de RMSE que generalmente se mantuvieron pequeñas, lo que indica que KANLib soporta un diseño arquitectónico flexible sin una pérdida significativa de precisión.

Significado y Reivindicaciones

El artículo posiciona a KANLib no como un modelo que logra la precisión predictiva de vanguardia sobre todos los métodos existentes, sino como una base robusta para la investigación futura. Su principal significado radica en:

• Cerrar la Brecha: Conecta el potencial teórico de las KAN con la aplicación práctica y escalable al unificar las implementaciones fragmentadas.
• Permitir la Extensibilidad: Al proporcionar una arquitectura modular, reduce la barrera para que los investigadores prueben nuevas funciones de base, estrategias de rejilla y arquitecturas híbridas.
• Estandarización: Ofrece un entorno consistente para evaluar las KAN, abordando el problema del soporte inconsistente de características que actualmente dificulta los estudios comparativos.

Los autores concluyen que, si bien KANLib incurre actualmente en una ligera sobrecarga computacional en comparación con el altamente especializado FastKAN (específicamente en el cálculo de GRBF debido a la inclusión de la lógica de rejilla adaptativa), proporciona la flexibilidad necesaria para explorar todo el espectro de las arquitecturas KAN. El trabajo futuro se identifica como la optimización del cálculo de la RBF gaussiana para cerrar la brecha de velocidad y la extensión del soporte para convoluciones 1D para datos de sensores dependientes del tiempo.