MAGE: Color-Invariant and Spatial Knowledge Distillation for Gastric Neoplasm Classification

Resumo Técnico: MAGE – Destilação de Conhecimento Invariante à Cor e Espacial para Classificação de Neoplasia Gástrica

Definição do Problema
A diferenciação precisa entre adenoma gástrico e carcinoma durante a endoscopia é crítica para a tomada de decisão clínica, mas permanece altamente desafiadora. A tarefa é dificultada pela alta similaridade interclasse, fronteiras ambíguas e a tendência das redes neurais profundas de depender de correlações espúrias (ex: sangramento, artefatos de iluminação ou contexto de fundo) em vez de características morfológicas intrínsecas. As abordagens existentes enfrentam limitações específicas:

• Métodos baseados em ROI: Frequentemente sofrem com propagação de erro devido à detecção ou segmentação imperfeita e perdem o contexto global circundante.
• Classificação de imagem completa: Carece do foco espacial necessário e frequentemente prediz com base em regiões de fundo em vez da própria lesão.
• Dependência de cor: Os modelos frequentemente gravitam para vieses de textura específicos de domínio, onde a cor atua como um confundidor frágil em vez de um indicador diagnóstico primário.

Metodologia: Estrutura MAGE
Os autores propõem o MAGE (Masked Achromatic Guidance Expert), uma estrutura de treinamento de dois ramos projetada para impor o aprendizado morfológico invariante à cor, mantendo a capacidade de implementação em imagens completas sem a necessidade de máscaras anotadas no momento da inferência.

1. Especialista Acromático Mascarado (Ramo Auxiliar):

   • Durante o treinamento, um ramo "especialista" auxiliar processa uma visão acromática mascarada da entrada.
   • A entrada é construída convertendo a Imagem de Luz Branca (WLI) para tons de cinza e aplicando uma máscara de segmentação de verdade fundamental ($\tilde{x} = m \odot \text{Gray}(x)$), definindo os pixels fora da Região de Interesse (ROI) como zero.
   • Este ramo é treinado exclusivamente em características estruturais e invariantes à cor, forçando-o a aprender padrões discriminativos baseados na morfologia, em vez de pistas de cor ou artefatos de fundo.
   • Este ramo utiliza informações privilegiadas (máscaras) e é congelado durante o processo de destilação.

2. Modelo Principal (Ramo Implementável):

   • O ramo principal recebe a WLI completa e bruta (RGB) como entrada e é o único componente utilizado durante a inferência.
   • Ele aprende a imitar o especialista por meio de uma estratégia de destilação de objetivos duplos:
     • Destilação de Logit de Classificação: Os logits de saída do modelo principal são regularizados via divergência KL contra os logits do especialista, transferindo a confiança de classificação do especialista.
     • Transferência de Atenção Espacial: O modelo principal é regularizado para alinhar seu mapa de atenção de características com o do especialista. Ao colapsar a dimensão do canal e calcular a distância de Frobenius ao quadrado entre o mapa de atenção do modelo principal e o do especialista (que é restrito à ROI), o modelo principal é implicitamente guiado a focar na estrutura da neoplasia dentro da imagem completa.

3. Escolha de Arquitetura:

   • A estrutura utiliza um backbone de Rede Neural Convolucional (CNN) em vez de Vision Transformers (ViTs). Os autores argumentam que as CNNs fornecem os vieses indutivos necessários (invariância à translação, localidade) adequados ao tamanho do conjunto de dados e garantem o alinhamento espacial preciso exigido para o cálculo da perda de atenção.

4. Função de Perda:
   O erro total $L$ combina:

   • Entropia Cruzada Binária (BCE) ponderada contra os rótulos de verdade fundamental.
   • Divergência KL escalonada por temperatura entre o especialista e o modelo principal.
   • Termo de regularização de atenção (distância de Frobenius entre mapas de atenção).

Principais Contribuições

1. Especialista Auxiliar de Treinamento Único: Introdução de um especialista acromático mascarado que promove o aprendizado morfológico invariante à cor ao suprimir o contexto de fundo e a informação de cor durante o treinamento.
2. Estrutura de Destilação Espacial: Um mecanismo para transferir pistas estruturais centradas na ROI e atenção espacial do especialista para um classificador de imagem completa implementável, permitindo a inferência em tempo real sem exigir máscaras de segmentação.
3. Desempenho e Robustez: Demonstração de melhorias consistentes sobre bases de detecção fortes (ex: YOLO) e de classificação (ex: Swin-Transformer, GastroNet), juntamente com maior explicabilidade e robustez a perturbações de cor.

Resultos Experimentais
Os experimentos foram conduzidos em um conjunto de dados privado de endoscopia gástrica (Seoul St. Mary's Hospital) e um conjunto de dados externo de cólon (PICCOLO).

• Desempenho de Classificação: O MAGE alcançou um AUC de 0,786 e F1 de 0,733 no teste gástrico, superando as bases, incluindo GastroNet-5M (AUC 0,744), Swin-Tiny (AUC 0,719) e pipelines baseados em YOLO.
• Explicabilidade: O MAGE demonstrou um alinhamento superior com as máscaras de verdade fundamental nas visualizações Grad-CAM.
  • Atenção na Máscara (AiM): MAGE alcançou 0,5303 (vs. 0,3176 para a base EffNet).
  • Precisão do Jogo de Apontar (PGA): MAGE alcançou 0,7324 (vs. 0,4355 para a base).
  • Isso indica que o modelo depende mais da morfologia da neoplasia do que de padrões globais espúrios.
• Robustez à Cor: Sob mudanças sintéticas de matiz, o MAGE mostrou uma degradação de desempenho significativamente menor em comparação com as bases EfficientNet e ResNet50 pré-treinada em GastroNet, confirmando a eficácia do treinamento invariante à cor.
• Generalização Cross-Domain: No conjunto de dados externo PICCOLO, o MAGE alcançou o maior AUC (0,9434), sugerindo que as características morfológicas aprendidas são transferíveis entre diferentes domínios anatômicos.
• Estudos de Ablação: Os resultados confirmaram que o design assimétrico (Especialista em Tons de Cinza + Principal em RGB) é crítico. O especialista teve melhor desempenho em entradas de tons de cinza, enquanto o modelo principal exigiu entrada RGB completa para otimizar a implementação.

Significância e Alegações
O artigo afirma que o MAGE aborda o problema do "aprendizado de atalho" (shortcut learning) na classificação de neoplasia gástrica ao forçar explicitamente o modelo a aprender características estruturais e invariantes à cor através de um sinal de treinamento privilegiado. Ao desacoplar o requisito de máscara do treinamento do requisito de inferência, o MAGE oferece uma solução prática que melhora tanto a precisão diagnóstica quanto a interpretabilidade clínica. Os autores enfatizam que seu método fornece não apenas métricas de classificação superiores, mas também mapas de atenção interpretáveis, que são essenciais para a confiabilidade clínica.

Limitações
Os autores observam que o estudo foca em classificação e robustez sem treinar um módulo de detecção/segmentação dedicado, o que pode limitar saídas de localização direta em certos fluxos de trabalho clínicos. Sugere-se que trabalhos futuros utilizem o classificador aprendido para localização leve (ex: baseado em CAM) e realizem uma validação multicêntrica mais ampla.