Multi-Dataset Cross-Domain Knowledge Distillation for Unified Medical Image Segmentation, Classification, and Detection

Resumo Técnico: Destilação de Conhecimento Cross-Domain Multi-Dataset para Segmentação, Classificação e Detecção Unificadas de Imagens Médicas

Declaração do Problema
A aprendizagem profunda revolucionou a análise de imagens médicas, contudo, modelos treinados em conjuntos de dados únicos frequentemente falham em generalizar para novos ambientes clínicos devido a deslocamentos de domínio causados por variações em protocolos de imagem, hardware de scanners e populações de pacientes. Embora a aprendizagem por transferência e a augmentação de dados ofereçam soluções parciais, elas frequentemente dependem de pré-treinamento em fonte única ou dados sintéticos que falham em capturar detalhes anatômicos de alta granularidade e variabilidade inter-dataset. Além disso, as abordagens existentes de destilação de conhecimento (KD) tipicamente utilizam um único professor em dados homogêneos para uma única tarefa (por exemplo, segmentação), negligenciando a natureza diversa e multi-tarefa dos fluxos de trabalho clínicos, que abrangem segmentação, classificação e detecção de objetos através de modalidades heterogêneas (Ressonância Magnética, Tomografia Computadorizada, Raio-X).

Metodologia
Os autores propõem um framework unificado e agnóstico à tarefa, baseado em professor-aluno, projetado para agregar conhecimento de múltiplos conjuntos de dados de imagens médicas heterogêneos. O framework opera em três estágios sequenciais:

1. Treinamento de Professor Adversarial de Domínio:

   • Um professor-alvo é treinado exclusivamente no conjunto de dados-alvo usando uma perda específica da tarefa (por exemplo, perda Dice para segmentação).
   • Professores-fonte são treinados em seus respectivos conjuntos de dados-fonte, mas aumentados com um discriminador de domínio. Utilizando um objetivo adversarial de domínio, esses codificadores são forçados a produzir características invariantes ao domínio que confundem o discriminador, alinhando as distribuições de características da fonte e do alvo. Este estágio suporta tanto adaptação supervisionada (usando dados-alvo rotulados) quanto não supervisionada.

2. Construção de Professor Conjunto via Fusão por Atenção Cruzada:

   • Um professor conjunto ($\mathcal{T}^*$) é construído fundindo as características de codificador e gargalo de múltiplos níveis do professor-alvo e de todos os professores-fonte.
   • Esta fusão utiliza mecanismos de atenção cruzada para agregar representações complementares de diferentes domínios. As camadas de codificador e gargalo de todos os professores constituintes são congeladas durante este processo para preservar o conhecimento aprendido.
   • Uma nova cabeça específica da tarefa (decodificador, classificador ou cabeça de detecção) é acoplada ao backbone fundido e treinada do zero no conjunto de dados-alvo. Isso permite que o professor conjunto se especialize na tarefa-alvo, mantendo características invariantes ao domínio.

3. Destilação Multi-nível Orientada por Currículo:

   • Um modelo estudante compacto é treinado no conjunto de dados-alvo para herdar conhecimento do professor conjunto.
   • O objetivo de treinamento combina uma perda de tarefa primária com múltiplos termos de destilação: alinhamento de características (MSE), similaridade cosseno, perda contrastiva (divergência KL em distribuições de similaridade) e extensões específicas da tarefa (por exemplo, transferência de atenção, destilação de logits, alinhamento de características de ROI para detecção).
   • Um fator de currículo ($c_f$) agenda dinamicamente o peso das perdas de destilação. O treinamento começa com foco na perda de tarefa supervisionada, aumentando gradualmente a influência dos termos de destilação para refinar as representações internas do estudante sem desestabilizar a aprendizagem inicial.

O framework é projetado para ser flexível, suportando arquiteturas UNet, TResUNet, MedViT, EfficientNet, Faster R-CNN e RF-DETR em tarefas de segmentação, classificação e detecção.

Principais Contribuições

1. Framework Unificado: Uma arquitetura única e escalável capaz de lidar com segmentação, classificação em nível de imagem e detecção de objetos através de conjuntos de dados médicos heterogêneos.
2. Arquitetura de Professor Conjunto: Um método inovador para fundir características de múltiplos níveis de múltiplos professores-fonte via atenção cruzada, permitindo transferência robusta de representação cross-domain sem exigir um único modelo pré-treinado monolítico.
3. Destilação Baseada em Currículo: Uma estratégia de destilação multi-nível que combina perdas de segmentação, contrastivas, de alinhamento e de similaridade, governada por um cronograma de currículo para garantir otimização estável.
4. Avaliação Abrangente: Experimentos extensivos através de três tarefas, três modalidades (Ressonância Magnética, Tomografia Computadorizada, Raio-X) e numerosos conjuntos de dados (incluindo BrainMetShare, ISLES, BraTS, LiTS, KiTS, COVIDx-CXR, OASIS, ADNI e DeepLesion).

Resultados Experimentais
O método proposto foi avaliado contra duas linhas de base: modelos específicos de conjunto de dados (treinados independentemente em cada conjunto de dados) e modelos multi-cabeça multi-conjunto de dados (backbone compartilhado com cabeças específicas de conjunto de dados).

• Segmentação: Em conjuntos de dados de Ressonância Magnética e Tomografia Computadorizada, os modelos estudantes superaram consistentemente ambas as linhas de base. Para tarefas baseadas em Tomografia Computadorizada, o estudante alcançou melhorias médias de IoU de ~2,5% sobre as linhas de base específicas de conjunto de dados e ~4,8% sobre modelos multi-cabeça, com reduções significativas na Distância de Hausdorff (HD95), indicando previsões de limites mais nítidas.
• Classificação: Em tarefas de classificação de doenças pulmonares e demência, os modelos estudantes destilados alcançaram maior precisão do que as linhas de base específicas de conjunto de dados e multi-cabeça. A destilação late-$k$ (transferência de conhecimento de camadas mais profundas) geralmente produziu resultados superiores.
• Detecção de Objetos: Em conjuntos de dados de Tomografia Computadorizada de pulmão, o framework melhorou a Precisão Média (mAP). O professor conjunto frequentemente alcançou o desempenho mais alto, enquanto o estudante destilado manteve a maior parte desses ganhos, oferecendo um modelo mais compacto adequado para implantação.
• Estudos de Ablação: Os resultados confirmaram que (i) a adaptação supervisionada ao alvo melhora a precisão dos limites; (ii) o desempenho escala com o número de conjuntos de dados-fonte; e (iii) a extração de características multi-nível (codificador + gargalo) produz melhores resultados do que a extração de nível único.
• Visualização: Visualizações t-SNE demonstraram que os modelos estudantes destilados aprendem clusters de características mais compactos e melhor separados em comparação com as linhas de base específicas de conjunto de dados em todas as três tarefas.

Significado e Alegações
O artigo afirma que a destilação de conhecimento multi-dataset serve como uma abordagem escalável e agnóstica à tarefa para aprimorar o desempenho através de domínios heterogêneos de imagens médicas. Ao agregar representações invariantes ao domínio de diversas fontes e destilá-las em um estudante compacto, o framework aborda os desafios de escassez de dados, custos de anotação e deslocamentos de domínio. Os autores enfatizam que sua abordagem produz melhorias consistentes sobre linhas de base fortes, demonstrando robustez aprimorada a deslocamentos distribucionais e generalização superior. O trabalho posiciona-se como uma extensão substancial de pesquisas anteriores focadas em segmentação, agora validada em todo o espectro de tarefas de análise de imagens médicas.