MedFormer-UR: Uncertainty-Routed Transformer for Medical Image Classification

Resumo Técnico: MedFormer-UR

1. Problema e Motivação

A integração segura de modelos de aprendizado profundo na prática clínica exige mais do que alta precisão; é fundamental a quantificação confiável da incerteza. Embora os Vision Transformers (ViTs) médicos, como o MedFormer, apresentem desempenho superior em tarefas de classificação e segmentação, eles enfrentam desafios críticos:

• Superconfiança: Modelos determinísticos frequentemente produzem probabilidades mal calibradas, exibindo alta confiança em entradas ambíguas ou fora da distribuição (OOD).
• Caixa Preta: A falta de transparência e interpretabilidade nas decisões dificulta a adoção clínica.
• Dados Ruídos e Desequilibrados: A natureza ruidosa e desbalanceada dos dados clínicos agrava a instabilidade das previsões.

O objetivo deste trabalho é desenvolver um framework unificado que não apenas classifique imagens médicas com precisão, mas que também quantifique a incerteza, localize ambiguidades e forneça decisões interpretáveis.

2. Metodologia

O artigo propõe o MedFormer-UR, uma extensão do MedFormer (um Transformer hierárquico) que integra três mecanismos principais para criar um sistema consciente da incerteza e baseado em protótipos:

A. Estimativa de Incerteza Evidencial por Token

• Em vez de usar apenas softmax determinístico, o modelo utiliza uma distribuição Dirichlet para modelar a incerteza evidencial.
• Um cabeçalho leve (MLP) prevê valores de evidência não negativos para cada token (região da imagem) e cada classe.
• A incerteza é derivada da força da distribuição Dirichlet: baixa evidência resulta em alta incerteza. Isso permite a quantificação e localização espacial da ambiguidade em tempo real (separando incerteza aleatória e epistêmica).

B. Roteamento Guiado por Incerteza (Uncertainty-Guided Routing)

• A incerteza não é apenas um output, mas um sinal de controle ativo durante o treinamento e inferência.
• O modelo emprega um mecanismo de Refinamento de Tokens Guiado por Incerteza (UGTR).
• Mecanismo: Em cada bloco do Transformer, há dois caminhos: um caminho principal de atenção (DSSA) e um ramo de refinamento (focado em detalhes locais).
• Uma máscara de roteamento suave é gerada para decidir onde aplicar o refinamento. Tokens com alta incerteza são suprimidos ou refinados de forma conservadora, evitando que atualizações de características não confiáveis degradem o aprendizado.
• Isso garante que o modelo foque em regiões informativas e seja cauteloso em áreas ambíguas.

C. Aprendizado Baseado em Protótipos

• Para melhorar a interpretabilidade, o modelo utiliza uma cabeça de classificação baseada em protótipos em vez de uma camada linear padrão.
• O modelo aprende vetores de referência (protótipos) específicos para cada classe.
• A classificação é feita medindo a similaridade (ex: cosseno) entre as características da imagem e esses protótipos.
• Isso estrutura o espaço de incorporação, permitindo que as decisões sejam baseadas em similaridade visual com padrões de referência conhecidos, facilitando a explicação clínica ("a imagem se assemelha a este protótipo de tumor maligno").

Função de Objetivo:
O treinamento otimiza uma função composta que inclui:

1. Perda de Entropia Cruzada (Classificação).
2. Perda de Supervisão de Roteamento (quando máscaras de tecido estão disponíveis).
3. Regularização de Protótipos (perda de cluster para agrupar tokens similares e perda de diversidade para evitar colapso de protótipos).

3. Contribuições Principais

1. Extensão do MedFormer: Introdução de cabeças evidenciais que produzem parâmetros Dirichlet para decomposição explícita da incerteza.
2. Mecanismo de Roteamento Ativo: Uso da incerteza estimada para modular dinamicamente a contribuição dos tokens, filtrando características ambíguas durante o treinamento.
3. Estruturação do Espaço de Embedding: Integração de protótipos aprendíveis com perda de contraste, permitindo raciocínio baseado em similaridade e previsões interpretáveis.
4. Avaliação Abrangente: Validação em quatro modalidades distintas (mamografia, ultrassom, ressonância magnética e histopatologia), demonstrando melhorias consistentes na calibração e seleção de previsões.

4. Resultados Experimentais

O modelo foi testado em quatro conjuntos de dados públicos: CBIS-DDSM (mamografia), BUSI (ultrassom), IDC (histopatologia) e Brain MRI (tumores cerebrais).

• Desempenho de Classificação:

  • O modelo proposto superou o MedFormer baseline em todas as modalidades.
  • Destaque no conjunto BUSI (ultrassom), onde a precisão aumentou de 0.800 para 0.893 e o AUROC de 0.919 para 0.956.
  • No CBIS-DDSM (mamografia), tarefa notoriamente difícil devido ao baixo contraste, houve um aumento na precisão de 0.629 para 0.673 e AUROC de 0.727 para 0.755.

• Calibração e Incerteza (Ponto Crítico):

  • A métrica chave foi o Erro de Calibração Esperado (ECE). O MedFormer baseline apresentou ECE alto (0.1836), indicando superconfiança.
  • O MedFormer-UR reduziu o ECE em até 35% (atingindo 0.0688 com parâmetro de roteamento $\beta=0.8$).
  • Isso significa que quando o modelo diz ter 80% de confiança, a probabilidade real de acerto está muito mais próxima desse valor.

• Previsão Seletiva:

  • O modelo demonstrou capacidade superior de rejeitar previsões incertas. Ao rejeitar 30% das previsões mais incertas (cobertura de 70%), a precisão do modelo aumentou significativamente, superando o baseline em todos os níveis de cobertura.

• Análise de Parâmetro $\beta$:

  • O estudo mostrou que a força do roteamento guiado por incerteza ($\beta$) não é monotônica. Valores intermediários (entre 0.7 e 0.8) ofereceram o melhor equilíbrio entre discriminação e confiabilidade calibrada. Valores extremos ($\beta=1.0$) melhoraram a precisão bruta, mas degradaram a calibração.

5. Significado e Conclusão

O MedFormer-UR representa um avanço significativo para a IA Confiável (Trustworthy AI) em medicina.

• Segurança Clínica: Ao reduzir a superconfiança e fornecer mapas de incerteza que correspondem a regiões clinicamente ambíguas (como bordas de lesões ou tecido denso), o modelo auxilia radiologistas a identificar casos que exigem atenção humana.
• Interpretabilidade: A combinação de mapas de incerteza e similaridade com protótipos oferece uma explicação transparente de "por que" uma decisão foi tomada, alinhando-se melhor com o fluxo de trabalho clínico do que as "caixas pretas" tradicionais.
• Generalização: A abordagem demonstra robustez em diferentes modalidades de imagem, sugerindo que a gestão ativa da incerteza é uma estratégia fundamental para a implantação segura de modelos de IA em ambientes de saúde reais.

Em suma, o trabalho prova que é possível melhorar não apenas a precisão, mas principalmente a confiabilidade e a transparência dos modelos de diagnóstico por imagem, tornando-os parceiros mais seguros para os profissionais de saúde.