Predicting Functions, Not Features: KANs with Function-Space Joint-Embedding Predictive Learning for Medical Image Segmentation

Resumo Técnico: Prevendo Funções, Não Características – KANs com Aprendizado Preditivo de Incorporação Conjunta no Espaço de Funções

Declaração do Problema

A segmentação de imagens médicas enfrenta desafios como contraste fraco, limites ambíguos e grandes variações de aparência. Embora arquiteturas de codificador-decodificador (ex: variantes de U-Net) abordem isso por meio de fusão multiescala, elas dependem de pesos escalares e ativações fixas, deixando as transformações não lineares definidas implicitamente.

As Redes Kolmogorov–Arnold (KANs) oferecem uma representação funcional explícita ao parametrizar as arestas da rede como funções univariadas aprendíveis ($\phi_{j,i}$). No entanto, os modelos de segmentação baseados em KAN existentes (ex: U-KAN) otimizam apenas as funções das arestas através de objetivos definidos após a agregação das arestas. Isso cria duas limitações primárias:

1. Falta de Alvos Explícitos Pré-Agregação: A supervisão restringe a soma das respostas das arestas de entrada, permitindo que diferentes configurações de funções de aresta produzam saídas de nós idênticas. Consequentemente, as funções de aresta individuais carecem de um alvo de aprendizado explícito antes da agregação.
2. Caracterização Insuficiente do Comportamento Local: Durante uma passagem direta (forward pass), uma função é avaliada apenas em um único âncora dependente da entrada. Funções com respostas idênticas nesse ponto podem exibir comportamentos locais diferentes nas proximidades, mas essas variações permanecem implicitamente restringidas pela tarefa de jusante (downstream).

Métodos de Aprendizado Preditivo de Incorporação Conjunta (JEPA) existentes operam em incorporações de características, tokens ou imagens, mas não tratam as funções neurais aprendíveis como alvos de predição. Estender o JEPA para o espaço de funções KAN não é trivial porque uma única resposta de aresta não pode caracterizar o comportamento local, e as arestas amostradas devem permanecer explicitamente alinhadas entre os ramos.

A metodologia proposta é o FS-JEPA (Aprendizado Preditivo de Incorporação Conjunta no Espaço de Funções), um framework que desloca o aprendizado preditivo para o espaço de funções pré-agregação das KANs.

Estrutura Central

A arquitetura consiste em uma rede de segmentação KAN em formato de U padrão e um ramo preditivo utilizado apenas durante o treinamento.

• Ramo Alvo (Target Branch): Um ramo de Média Móvel Exponencial (EMA) de contexto total gera assinaturas alvo.
• Ramo Online (Online Branch): Um ramo mascarado recebe uma sequência de tokens parcialmente mascarada e prevê as assinaturas alvo.
• Correspondência de Arestas: Ambos os ramos compartilham os mesmos índices de arestas amostrados para garantir que as predições correspondam às mesmas funções de aresta de entrada-saída.

Componentes Principras

1. Amostragem e Seleção de Arestas
Para evitar o custo proibitivo de computar as respostas para todas as $O(BN d_{in} d_{out})$ arestas, o framework amostra aleatoriamente $K$ pares de arestas de entrada-saída. Um operador de aresta selecionada ($\Phi_S$) avalia apenas essas arestas específicas coletando seus parâmetros de spline e grades de nós (knot grids), reduzindo os custos de ativação auxiliar para $O(BNK)$.

2. Assinatura de Função de Aresta de Multi-Raio
Em vez de prever uma única resposta isolada, o método representa cada função de aresta amostrada usando uma assinatura de multi-raio.

• Para uma aresta com âncora de entrada $u$, um conjunto de deslocamentos locais $\mathcal{D} = \{-0.10, -0.05, -0.025, 0.025, 0.05, 0.10\}$ é definido.
• A assinatura $s_e$ é o vetor das avaliações da função nesses deslocamentos: $s_e = [\phi_e(u+\Delta)]_{\Delta \in \mathcal{D}}$.
• Esta representação estruturada captura variações funcionais locais que uma avaliação de ponto único não consegue capturar.

3. Correspondência de Arestas Determinística
Para resolver o problema da amostragem estocástica, as posições de origem e destino dos nós de uma aresta amostrada $(i, j)$ são codificadas como coordenadas determinísticas $(p_i, p_j)$. Essas coordenadas são concatenadas com a assinatura da função e fornecidas como condições para o preditor de assinatura, garantindo consistência sem parâmetros aprendíveis.

4. Otimização Conjunta
O framework otimiza conjuntamente a perda de segmentação ($L_{seg}$) e a perda preditiva no espaço de funções ($L_{sig}$).

• Perda Preditiva: Definida como $L_{sig} = \text{SmoothL1}(\hat{s}_e, s^+_e)$, onde $\hat{s}_e$ é a predição do ramo online e $s^+_e$ é o alvo do ramo EMA.
• Treinamento: A perda total é $L_{total} = L_{seg} + \lambda_{deep}L_{deep} + \lambda_{sig}\sum L^{(b)}_{sig}$.
• Inferência: Todos os componentes preditivos (ramo online, ramo EMA, preditor de assinatura) são removidos, deixando o caminho de segmentação original inalterado.

Principais Contribuições

1. JEPA no Espaço de Funções: A formulação do aprendizado preditivo de incorporação conjunta no espaço de funções nativo de pré-agregação das KANs, deslocando o alvo de predição de incorporações de características para assinaturas estruturadas de funções de aresta.
2. Representação de Assinatura de Multi-Raio: A introdução de uma representação que captura o comportamento funcional local das funções de aresta KAN, fornecendo um alvo de predição eficaz para o aprendizado de funções.
3. Framework Preditivo de Uso Apenas no Treinamento: Um método que aprende conjuntamente assinaturas de funções de aresta KAN e representações de segmentação a partir de inicialização aleatória, mantendo a correspondência de arestas, sem custo adicional de parâmetros ou computação durante a inferência.

Resultados Experimentais

O método foi avaliado em cinco benchmarks de segmentação de imagens médicas: BUSI (ultrassonografia), DDTI (ultrassonografia), TN3K (tireoide), CVC-ClinicDB (pólipos) e GlaS (glândulas).

• Desempenho: O FS-JEPA alcançou o maior Dice médio (83,37%) e IoU (75,04%) nos cinco conjuntos de dados.
• Comparação: Superou o método baseado em KAN mais forte (UUEKAN) em +2,25 pontos percentuais no Dice.
• Ablação:
  • Comparado a um scaffold correspondente treinado sem FS-JEPA, o método proposto melhorou o Dice médio em 0,85 pontos percentuais.
  • A ablação sobre alvos preditivos mostrou que prever assinaturas de funções de aresta de multi-raio (0,7926 Dice) superou significativamente a predição de respostas de aresta única (0,7771 Dice) ou características de token/nó.
  • A análise mecanística confirmou que o FS-JEPA fornece um sinal de aprendizado de pré-agregação explícito que molda diretamente as representações das funções de aresta KAN, tornando-as sensíveis a variações no nível da aresta que a perda de segmentação padrão ignora.

Significância e Alegações

O artigo afirma que o FS-JEPA demonstra a eficácia da extensão do aprendizado preditivo além das convenções de incorporação de características para as representações estruturadas de funções neurais aprendíveis. Ao prever assinaturas de funções de aresta de multi-raio antes da agregação, o método permite o aprendizado de representação direta sobre as funções de aresta KAN. Os autores afirmam que as melhorias de desempenho originam-se dessa supervisão preditiva no espaço de funções, e não de um aumento na capacidade do modelo, visto que a contagem de parâmetros no tempo de inferência permanece idêntica ao scaffold de base. O framework oferece uma maneira de restringir explicitamente o comportamento funcional de arestas KAN individuais, abordando uma limitação fundamental nos modelos de visão baseados em KAN existentes.