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

Resumo Técnico de "TruKAN: Rumo a Redes Kolmogorov-Arnold Mais Eficientes Usando Funções de Potência Truncadas"

Problema

As Redes Kolmogorov-Arnold (KANs) surgiram como uma alternativa promissora aos Perceptrons de Camada Multicamadas (MLPs), oferecendo interpretabilidade superior e a capacidade de aprender funções univariadas ao longo das arestas. No entanto, as KANs padrão dependem de funções de base B-spline para suas ativações aprendíveis. A avaliação dessas B-splines requer o algoritmo recursivo de de Boor–Cox, que é computacionalmente caro, fazendo com que as KANs treinem até 10 vezes mais devagar que MLPs equivalentes. Além disso, a dependência de algoritmos de spline complexos e estruturas de grade fixas pode levar a problemas de escalabilidade em problemas de alta dimensão e arquiteturas profundas, caracterizados por gradientes ruidosos e otimização instável. Embora variantes recentes como SineKAN e FourierKAN tentem abordar esses problemas, elas frequentemente introduzem novas compensações em relação à interpretabilidade ou permanecem limitadas a benchmarks sintéticos e conjuntos de dados simples como o MNIST.

Metodologia

Os autores propõem a TruKAN, uma nova arquitetura que mantém a topologia canônica das KANs e as ativações aprendíveis, mas substitui a base B-spline por uma família de funções de potência truncadas derivadas da teoria de splines de ordem $k$.

Componentes Arquiteturais Principais:

1. Base de Potência Truncada: Em vez de B-splines, a TruKAN utiliza funções de potência truncada da forma $(x - t_j)^k_+$. Teoricamente, o espaço gerado por essas funções é idêntico ao das B-splines, mas sua avaliação computacional é distinta. Enquanto as B-splines possuem suporte compacto (exigindo complexidade $O(k^2)$ por intervalo via recursão), as funções de potência truncada possuem suporte global. No entanto, os autores argumentam que, na prática, grades de alta resolução não são sempre necessárias, permitindo que a TruKAN opere eficientemente com um pequeno número de nós (knots).
2. Estrutura de Camada Híbrida: Cada camada da TruKAN combina dois componentes:
   • Um componente de potência truncada para capturar não-linearidades locais e por partes.
   • Um componente polinomial (tendência global) para substituir a ativação auxiliar SiLU usada em KANs padrão. Este termo polinomial fornece um mapeamento global suave e diferenciável, compensando a falta de uma não-linearidade global explícita como a SiLU.
3. Configurações de Nós (Knots): O artigo introduz duas variantes em relação ao posicionamento dos nós:
   • Nós Compartilhados (Shared Knots): Um único conjunto de localizações de nós é compartilhado entre todos os canais de saída.
   • Nós Individuais (Individual Knots): Cada canal de saída possui seu próprio conjunto de posições de nós aprendíveis.
   • Estratégias de nós tanto fixas (intervalos pré-definidos) quanto aprendíveis (otimizadas durante o treinamento com restrições de ordenação e positividade) são exploradas.
4. Estrutura de Treinamento:
   • Otimização: Para abordar a instabilidade do treinamento de redes profundas com funções de potência truncada (que podem sofrer explosão de gradiente devido ao crescimento/decaimento exponencial), os autores empregam uma estratégia de otimização híbrida usando AdamW envolto por LookAhead. Isso substitui o L-BFGS padrão usado em estudos anteriores de KAN para melhor se adequar ao treinamento estocástico de mini-lotes (mini-batch).
   • Normalização: O estudo investiga o impacto da Normalização de Camada (Layer Normalization) (e Batch Normalization) dentro dos cabeçalhos de classificação, já que técnicas de regularização tradicionais como dropout mostraram resultados mistos em KANs.
   • Backbone: Os modelos são integrados a uma estrutura EfficientNet-V2, substituindo o cabeçalho de classificador padrão por camadas MLP, KAN, SineKAN ou TruKAN.

Contribuições Principais

• Inovação Arquitetural: A substituição das B-splines por funções de potência truncada, juntamente com um termo global polinomial, cria um modelo que mantém a expressividade da KAN enquanto simplifica a avaliação da função de base.
• Interpretabilidade Aprimorada: Ao usar funções de potência truncada, a contribuição de cada nó torna-se explícita. Um coeficiente não nulo em um nó $t_j$ indica diretamente uma curvatura de ordem $k$ intencional naquele ponto específico, oferecendo maior transparência do que o suporte sobreposto das B-splines.
• Escalabilidade e Eficiência: O artigo demonstra que a TruKAN pode ser efetivamente escalada para arquiteturas profundas e tarefas de visão complexas, indo além das limitações de dados sintéticos de estudos anteriores de KAN.
• Benchmarking Abrangente: Os autores fornecem uma avaliação rigorosa em quatro conjuntos de dados de visão computacional (CIFAR-10, CIFAR-100, Oxford-II-Pets, STL-10) usando configurações de rede pequenas e profundas, garantindo comparações justas ao igualar a contagem de parâmetros.

Resultados Experimentais

A avaliação foi conduzida em quatro conjuntos de dados usando backbones EfficientNet-V2 com tamanhos de classificador variados (arquiteturas pequenas e grandes).

• Acurácia e F1-Score:
  • Arquiteturas Pequenas: As variantes TruKAN, particularmente as com Normalização de Camada e Nós Compartilhados (TruKAN-SN), superaram consistentemente ou igualaram os modelos de linha de base KAN e SineKAN. No CIFAR-10, a TruKAN-IN alcançou a maior acurácia (0.8927), superando o melhor baseline MLP. No STL-10, a TruKAN-SN alcançou 0.9341 de acurácia, ficando atrás apenas do MLP normalizado por uma margem insignificante.
  • Arquiteturas Grandes: Os modelos TruKAN demonstraram escalabilidade superior. No CIFAR-10, a TruKAN-IN alcançou 0.8932 de acurácia, superando o baseline MLP-N apesar de usar menos parâmetros. As variantes TruKAN dominaram consistentemente os baselines KAN e SineKAN em todos os datasets, muitas vezes alcançando resultados comparáveis ou melhores que os MLPs mais fortes.
  • Impacto da Normalização: A normalização de camada provou ser crítica para o desempenho da TruKAN, melhorando significativamente a generalização, particularmente para configurações de nós individuais, que de outra forma sofreriam overfitting.
• Eficiência Computacional:
  • Tempo de Treinamento: Modelos TruKAN de nós compartilhados (TruKAN-FS) alcançaram tempos de passo de treinamento de ~0.0024s, representando uma melhoria de 3-4x sobre os modelos KAN padrão (0.0087s - 0.0099s).
  • Uso de Memória: A TruKAN-FS consumiu apenas 118.04 MB de memória de pico da GPU, significativamente menos que o KAN-PBT (570.84 MB) e o KAN-PBF (691.79 MB). No entanto, as variantes de nós individuais (TruKAN-FI) incorreram em altos custos de memória (~2.6 GB) devido à expansão do armazenamento dos nós.
  • Parâmetros: Os modelos TruKAN geralmente alcançaram desempenho comparável ou melhor com menos parâmetros treináveis do que seus equivalentes KAN e SineKAN.

Significância e Alegações

O artigo afirma que a TruKAN aborda com sucesso o equilíbrio entre eficiência computacional e adesão aos princípios das KANs. Ao substituir o algoritmo recursivo de B-spline por funções de potência truncada e um termo polinomial global, a TruKAN alcança:

1. Eficiência Superior: Reduções significativas no tempo de treinamento e uso de memória em comparação com as KANs padrão, tornando-as viáveis para tarefas de visão complexas.
2. Expressividade Mantida: O modelo preserva as capacidades de aproximação teórica das KANs enquanto melhora a acurácia em datasets do mundo real.
3. Interpretabilidade Melhorada: A natureza explícita das funções de potência truncada permite uma inspeção mais clara das funções de aresta aprendidas e das contribuições dos nós em comparação com as B-splines ou métodos espectrais globais como o SineKAN.

Os autores concluem que a TruKAN oferece um caminho prático para as KANs, demonstrando vantagens além das configurações restritas de estudos anteriores. Eles reconhecem que, embora o método melhore a estabilidade, desafios permanecem em relação à integração de outras técnicas de regularização (como dropout) e ao potencial do treinamento de precisão mista para aumentar ainda mais a eficiência. O trabalho futuro propõe explorar essas áreas e aplicar a TruKAN em aprendizado por reforço e previsão de séries temporais.