TruKAN: Towards More Efficient Kolmogorov-Arnold Networks Using Truncated Power Functions

Resumen Técnico de "TruKAN: Hacia Redes Kolmogorov-Arnold más Eficientes mediante Funciones de Potencia Truncadas"

Planteamiento del Problema

Las Redes Kolmogorov-Arnold (KAN) han surgido como una alternativa prometedora a los Perceptrones Multicapa (MLP), ofreciendo una interpretabilidad superior y la capacidad de aprender funciones univariantes en los bordes. Sin embargo, las KAN estándar dependen de funciones de base B-spline para sus activaciones aprendibles. La evaluación de estas B-splines requiere el algoritmo recursivo de de Boor–Cox, el cual es computacionalmente costoso, causando que las KAN se entrenen hasta 10 veces más lento que los MLP equivalentes. Además, la dependencia de algoritmos de spline complejos y estructuras de rejilla fijas puede derivar en problemas de escalabilidad en problemas de alta dimensión y arquitecturas profundas, caracterizados por gradientes ruidosos y una optimización inestable. Si bien variantes recientes como SineKAN y FourierKAN intentan abordar estos problemas, a menudo introducen nuevos compromisos respecto a la interpretabilidad o permanecen limitadas a benchmarks sintéticos y conjuntos de datos simples como MNIST.

Metodología

Los autores proponen TruKAN, una nueva arquitectura que retiene la topología canónica de las KAN y las activaciones aprendibles, pero reemplaza la base B-spline por una familia de funciones de potencia truncadas derivadas de la teoría de splines de orden $k$.

Componentes Arquitectónicos Centrales:

1. Base de Potencia Truncada: En lugar de B-splines, TruKAN utiliza funciones de potencia truncadas de la forma $(x - t_j)^k_+$. Teóricamente, el espacio generado por estas funciones es idéntico al de las B-splines, pero su evaluación computacional es distinta. Mientras que las B-splines tienen soporte compacto (requiriendo una complejidad de $O(k^2)$ por intervalo vía recursión), las funciones de potencia truncadas tienen soporte global. Sin embargo, los autores argumentan que, en la práctica, no siempre son necesarias rejillas de alta resolución, lo que permite que TruKAN opere eficientemente con un número pequeño de nudos (knots).
2. Estructura de Capa Híbrida: Cada capa de TruKAN combina dos componentes:
   • Un componente de potencia truncada para capturar no linealidades locales y segmentadas.
   • Un componente polinómico (tendencia global) para reemplazar la activación auxiliar SiLU utilizada en las KAN estándar. Este término polinómico proporciona un mapeo global suave y diferenciable, compensando la falta de una no linealidad global explícita como SiLU.
3. Configuraciones de Nudos: El artículo introduce dos variantes respecto a la colocación de los nudos:
   • Nudos Compartidos (Shared Knots): Un único conjunto de ubicaciones de nudos se comparte entre todos los canales de salida.
   • Nudos Individuales (Individual Knots): Cada canal de salida tiene su propio conjunto de posiciones de nudos aprendibles.
   • Se exploran tanto las estrategias de nudos fijos (intervalos predefinidos) como las de nudos aprendibles (optimizados durante el entrenamiento con restricciones de orden y positividad).
4. Marco de Entrenamiento:
   • Optimización: Para abordar la inestabilidad del entrenamiento de redes profundas con funciones de potencia truncadas (que pueden sufrir explosión de gradientes debido al crecimiento/decimiento exponencial), los autores emplean una estrategia de optimización híbrida utilizando AdamW envuelto por LookAhead. Esto reemplaza al L-BFGS estándar utilizado en estudios previos de KAN para adaptarse mejor al entrenamiento estocástico por mini-lotes (mini-batch).
   • Normalización: El estudio investiga el impacto de la Normalización de Capa (Layer Normalization) (y la Normalización de Lote/Batch Normalization) dentro de los cabezales clasificadores, ya que las técnicas de regularización tradicionales como el dropout han mostrado resultados mixtos en las KAN.
   • Backbone: Los modelos se integran en un marco de EfficientNet-V2, reemplazando el cabezal clasificador estándar por capas MLP, KAN, SineKAN o TruKAN.

Contribuciones Clave

• Innovación Arquitectónica: La sustitución de las B-splines por funciones de potencia truncadas, junto con un término global polinómico, crea un modelo que mantiene la expresividad de las KAN mientras simplifica la evaluación de la función de base.
• Interpretabilidad Mejorada: Al utilizar funciones de potencia truncadas, la contribución de cada nudo se vuelve explícita. Un coeficiente no nulo en un nudo $t_j$ indica directamente una curvatura de orden $k$ intencional en ese punto específico, ofreciendo mayor transparencia que el soporte superpuesto de las B-splines.
• Escalabilidad y Eficiencia: El artículo demuestra que TruKAN puede escalarse efectivamente a arquitecturas profundas y tareas de visión complejas, yendo más allá de las limitaciones de datos sintéticos de los estudios previos de KAN.
• Benchmarking Exhaustivo: Los autores proporcionan una evaluación rigurosa a través de cuatro conjuntos de datos de visión computacional (CIFAR-10, CIFAR-100, Oxford-IIIT Pets, STL-10) utilizando configuraciones de red pequeñas y profundas, asegurando comparaciones justas al igualar el conteo de parámetros.

Resultados Experimentales

La evaluación se realizó en cuatro conjuntos de datos utilizando backbones EfficientNet-V2 con tamaños de clasificador variables (arquitecturas pequeñas y grandes).

• Precisión y F1-Score:
  • Arquitecturas Pequeñas: Las variantes de TruKAN, particularmente aquellas con Normalización de Capa y Nudos Compartidos (TruKAN-SN), superaron consistentemente o igualaron a los modelos base KAN y SineKAN. En CIFAR-10, TruKAN-IN alcanzó la precisión más alta (0.8927), superando al mejor baseline MLP. En STL-10, TruKAN-SN alcanzó una precisión de 0.9341, quedando solo por detrás del MLP normalizado por un margen insignificante.
  • Arquitecturas Grandes: Los modelos TruKAN demostraron una escalabilidad superior. En CIFAR-10, TruKAN-IN alcanzó una precisión de 0.8932, superando al baseline MLP-N a pesar de usar menos parámetros. Las variantes de TruKAN dominaron consistentemente los baselines KAN y SineKAN en todos los conjuntos de datos, logrando a menudo resultados comparables o mejores que los MLPs más fuertes.
  • Impacto de la Normalización: La normalización de capa resultó crítica para el rendimiento de TruKAN, mejorando significativamente la generalización, particularmente para las configuraciones de nudos individuales, que de otro modo sufrían de sobreajuste.
• Eficiencia Computacional:
  • Tiempo de Entrenamiento: Los modelos de TruKAN con nudos compartidos (TruKAN-FS) lograron tiempos de paso de entrenamiento de ~0.0024s, lo que representa una mejora de 3-4 veces respecto a los modelos KAN estándar (0.0087s - 0.0099s).
  • Uso de Memoria: TruKAN-FS consumió solo 118.04 MB de memoria GPU pico, significativamente menos que KAN-PBT (570.84 MB) y KAN-PBF (691.79 MB). Sin embargo, las variantes de nudos individuales (TruKAN-FI) incurrieron en altos costos de memoria (~2.6 GB) debido a la expansión del almacenamiento de los nudos.
  • Parámetros: Los modelos TruKAN generalmente lograron un rendimiento comparable o mejor con menos parámetros entrenables que sus contrapartes de KAN y SineKAN.

Significancia y Reivindicaciones

El artículo afirma que TruKAN aborda con éxito el compromiso entre la eficiencia computacional y la adherencia a los principios de las KAN. Al reemplazar el algoritmo recursivo de B-spline con funciones de potencia truncadas y un término global polinómico, TruKAN logra:

1. Eficiencia Superior: Reducciones significativas en el tiempo de entrenamiento y el uso de memoria en comparación con las KAN estándar, haciéndolas viables para tareas de visión complejas.
2. Mantenimiento de la Expresividad: El modelo preserva las capacidades de aproximación teórica de las KAN mientras mejora la precisión en conjuntos de datos del mundo real.
3. Interpretabilidad Mejorada: La naturaleza explícita de las funciones de potencia truncadas permite una inspección más clara de las funciones de borde aprendidas y las contribuciones de los nudos en comparación con las B-splines o los métodos espectrales globales como SineKAN.

Los autores concluyen que TruKAN ofrece un camino práctico hacia adelante para las KAN, demostrando ventajas más allá de los entornos restringidos de estudios previos. Reconocen que, si bien el método mejora la estabilidad, aún quedan desafíos respecto a la integración de otras técnicas de regularización (como el dropout) y el potencial del entrenamiento de precisión mixta para mejorar aún más la eficiencia. Se propone como trabajo futuro explorar estas áreas y aplicar TruKAN al aprendizaje por refuerzo y la predicción de series temporales.