DerMAE: Improving skin lesion classification through conditioned latent diffusion and MAE distillation

O artigo apresenta o DerMAE, um método que utiliza modelos de difusão condicionados para gerar imagens sintéticas e pré-treinamento MAE para melhorar a classificação de lesões de pele em cenários de desequilíbrio de classes, transferindo o conhecimento para modelos leves via destilação para viabilizar sua implementação em dispositivos móveis.

O essencial

Imagine que você é um médico especialista em pele, mas você só tem acesso a um livro de fotos de doenças de pele que está desequilibrado.

Nesse livro, há milhares de fotos de pintas comuns e inofensivas (benignas), mas apenas algumas poucas fotos de câncer de pele (maligno). Se você tentar ensinar um computador a aprender com esse livro, ele vai ficar "preguiçoso": vai achar que toda mancha é inofensiva, porque é isso que ele vê o tempo todo. Quando ele finalmente encontrar um caso de câncer, ele provavelmente vai errar, achando que é apenas uma pintinha comum. Isso é perigoso.

Os autores deste trabalho, do Brasil, criaram uma solução inteligente em três etapas para consertar isso e colocar um "médico digital" no celular de qualquer pessoa. Vamos entender como funciona usando analogias simples:

1. O "Chef de Cozinha" que Cria Pratos Novos (Geração Sintética)

O primeiro problema é a falta de fotos de câncer. A equipe usou uma tecnologia chamada Modelo de Difusão (o mesmo tipo de IA que cria imagens a partir de texto, como o DALL-E ou Midjourney).

• A Analogia: Imagine que o livro de receitas (o conjunto de dados) tem 90% de receitas de bolo de cenoura e apenas 10% de bolo de chocolate. O chef (a IA) precisa aprender a fazer os dois.
• A Solução: Em vez de apenas cortar e colar pedaços das fotos existentes (o que é como tentar fazer um bolo de chocolate misturando cenoura com chocolate), eles usaram um "Chef de IA" para criar receitas de bolo de chocolate do zero.
• O Truque: Eles ensinaram esse chef a criar fotos de lesões malignas (câncer) especificamente quando pediam, equilibrando a "prateleira" de fotos. Eles também usaram um "gosto refinado" (uma perda perceptiva) para garantir que as fotos criadas parecessem reais e não apenas borrões coloridos.

2. O "Estudante de Elite" que Aprende Sozinho (Pré-treinamento MAE)

Agora que eles tinham muitas fotos (reais + criadas), precisavam ensinar um computador a entendê-las. Eles usaram um modelo gigante chamado ViT-H (uma rede neural muito poderosa, mas pesada).

• A Analogia: Imagine um estudante universitário brilhante (o modelo ViT-H) que precisa aprender a identificar doenças. Mas ele não pode apenas olhar para as fotos e tentar memorizar.
• O Truque: Eles cobriram 75% da foto com um "papel de parede" (máscara) e pediram para o estudante adivinhar o que estava escondido. Para fazer isso, ele teve que aprender a entender a estrutura global da lesão, e não apenas detalhes pequenos.
• O Resultado: Esse estudante aprendeu muito bem, usando as fotos criadas pelo "Chef" para se tornar um especialista em entender a diferença entre uma pintinha e um câncer, mesmo sem ter visto muitos exemplos reais de câncer no início.

3. O "Mestre" que Ensina um "Estagiário" (Distilação de Conhecimento)

O problema é que esse "Estudante de Elite" (ViT-H) é gigante. Ele é como um supercomputador que cabe em um servidor de dados, mas não cabe no seu celular. Se tentássemos rodar isso no celular de um médico em uma zona rural, o celular travaria.

• A Analogia: Precisamos de um Estagiário (um modelo menor, como o ViT-B ou EfficientNet) que seja leve, rápido e rode no celular. Mas o estagiário é "burro" comparado ao mestre.
• A Solução: Eles usaram uma técnica chamada Distilação de Conhecimento. É como se o Mestre (o modelo gigante pré-treinado) passasse um curso intensivo para o Estagiário.
  • O Mestre não apenas diz "Isso é câncer".
  • Ele explica por que é câncer, mostrando as nuances e a probabilidade de ser maligno.
  • O Estagiário aprende a pensar como o Mestre, mas com um cérebro muito menor e mais rápido.

O Resultado Final

No final, eles conseguiram:

1. Resolver o desequilíbrio: Criaram fotos de câncer suficientes para ensinar a IA a não ter medo de errar.
2. Aprendizado robusto: A IA aprendeu a ver padrões complexos usando o método de "adivinhar o que está escondido".
3. Leveza para o mundo real: Transferiram a inteligência do "gigante" para um "pequeno" que cabe no celular.

Por que isso importa?
Isso significa que, no futuro, um médico generalista ou até um paciente em uma área remota poderá tirar uma foto de uma lesão na pele com o celular e receber uma análise precisa e rápida, ajudando a detectar câncer de pele mais cedo, mesmo sem ter um dermatologista especialista por perto. A IA aprendeu a ser justa (não ignorando os casos raros) e eficiente (rodando em qualquer aparelho).

Resumo técnico

Resumo Técnico: DERMAE

1. O Problema

A classificação de lesões cutâneas enfrenta dois desafios críticos que limitam a eficácia dos modelos de aprendizado profundo atuais:

• Desequilíbrio Severo de Classes: Os conjuntos de dados dermatológicos (como o HAM10000) são altamente desbalanceados, com aproximadamente 90% das amostras sendo lesões benignas e uma sub-representação significativa de casos malignos. Isso leva os modelos a aprenderem fronteiras de decisão enviesadas, favorecendo a classe majoritária.
• Limitação de Dados e Arquiteturas: Modelos baseados em Vision Transformers (ViT), que possuem capacidades poderosas de aprendizado de representação, dependem de grandes volumes de dados para generalização. A escassez de dados reais e o desbalanceamento impedem que esses modelos atinjam seu potencial em dermatologia.
• Restrições de Implantação: Soluções clínicas práticas, especialmente em dispositivos móveis ou de ponta (edge), exigem modelos leves, enquanto os modelos de alta performance (como ViT-Huge) são computacionalmente proibitivos para esses cenários.

2. Metodologia

O trabalho propõe um pipeline de três etapas, denominado DERMAE, que integra geração sintética, pré-treinamento auto-supervisionado e destilação de conhecimento:

A. Geração de Dados Sintéticos (Latent Diffusion)

• Utiliza-se um Modelo Probabilístico de Difusão Denoising (DDPM) para gerar imagens dermatológicas sintéticas.
• Estratégias de Geração:
  1. Geração Condicional: Inserção de embeddings de classe durante a inferência para equilibrar explicitamente as classes malignas e benignas.
  2. Geração Não Condicional: Geração de 600.000 amostras para modelar a distribuição geral de lesões.
• Função de Perda: Para melhorar a fidelidade perceptual e reduzir artefatos de baixa frequência comuns em funções de perda MSE, o modelo incorpora uma perda perceptual (baseada em redes pré-treinadas) além da perda de erro quadrático médio.
• Arquitetura: Backbone U-Net com 1000 passos de tempo.

B. Pré-treinamento Auto-supervisionado (MAE)

• Um modelo Masked Autoencoder (MAE) com backbone ViT-Huge (ViT-H/16) é pré-treinado utilizando os dados sintéticos gerados.
• Objetivo: O modelo aprende a reconstruir patches de imagem mascarados (75% de taxa de mascaramento), forçando-o a aprender representações latentes robustas e semanticamente significativas das estruturas globais das lesões, sem depender de rótulos supervisionados.
• Isso permite que o modelo aprenda características de classes sub-representadas (malignas) a partir do vasto conjunto de dados sintéticos.

C. Destilação de Conhecimento (Knowledge Distillation)

• Para viabilizar a implantação em dispositivos móveis, o conhecimento do grande modelo "professor" (ViT-H pré-treinado com MAE) é transferido para um modelo "aluno" menor.
• Arquitetura do Aluno: ViT-Base (ViT-B/16) ou EfficientNet-B0.
• Mecanismo: Utiliza-se uma estratégia de soft-targets, onde o aluno é treinado para imitar tanto a distribuição de probabilidade do professor quanto os rótulos reais (ground-truth). A função de perda combina Entropia Cruzada e Divergência KL.
• O aluno é ajustado (fine-tuned) usando uma combinação de dados reais e sintéticos.

3. Contribuições Principais

1. Abordagem Híbrida para Desequilíbrio: Integração bem-sucedida de modelos de difusão condicionados por classe para gerar dados sintéticos de alta qualidade, mitigando o desbalanceamento de classes em dermatologia.
2. Pré-treinamento com MAE em Dados Sintéticos: Demonstração de que o pré-treinamento auto-supervisionado em grandes volumes de dados sintéticos permite que ViTs aprendam representações robustas, superando a limitação de dados reais.
3. Eficiência para Dispositivos Móveis: Aplicação de destilação de conhecimento para transferir a performance de modelos pesados (ViT-H) para arquiteturas leves (ViT-B e EfficientNet), permitindo inferência eficiente em dispositivos com recursos limitados (ex: ~0.6 GFLOPs para EfficientNet-B0 vs ~55-60 GFLOPs para ViT-H).
4. Validação em Cenário Realista: O framework foi testado no dataset HAM10000, demonstrando ganhos significativos tanto na classificação multiclasse quanto binária (Benigno/Maligno).

4. Resultados Experimentais

Os experimentos foram realizados no dataset HAM10000 (10.000 imagens, 8 classes).

• Desempenho de Classificação:
  • O modelo MAE + ViT-H + Synth alcançou a melhor performance geral, com 89,15% de precisão (ACC) e 89,11% de F1-score na classificação categórica, superando significativamente as linhas de base (ViT-B baseline: ~67% ACC).
  • Na classificação binária (Maligno vs Benigno), o modelo alcançou 87,15% de ACC e 87,15% de F1-score.
  • A combinação de Destilação + Augmentação (usando EfficientNet-B0 como aluno) obteve 90,51% de ACC e 90,10% de F1-score na classificação binária, demonstrando que modelos leves podem atingir performance superior ao usar o conhecimento do professor.
• Comparação com Baselines: O método superou variações do EfficientNet treinados apenas com dados reais e outros métodos de geração sintética (como Derm-t2im).
• Qualidade Visual: A inclusão da perda perceptual e do condicionamento de classe resultou em imagens sintéticas com maior fidelidade visual e diversidade estrutural, conforme demonstrado na análise qualitativa (Figura 2).

5. Significado e Impacto

O trabalho DERMAE oferece uma solução viável para a implantação de sistemas de diagnóstico de câncer de pele em cenários clínicos reais, especialmente em locais com poucos especialistas ou recursos computacionais limitados.

• Viabilidade Clínica: Ao permitir que modelos leves (adequados para celulares) herdem o conhecimento de modelos grandes treinados em dados sintéticos, o sistema torna-se acessível para triagem inicial por médicos generalistas ou em áreas remotas.
• Mitigação de Viés: A geração de dados sintéticos condicionados ajuda a corrigir o viés histórico dos datasets dermatológicos, melhorando a detecção de casos malignos que são frequentemente negligenciados.
• Escalabilidade: A metodologia demonstra que é possível criar pipelines de IA robustos para medicina sem depender exclusivamente da coleta massiva de novos dados reais, contornando barreiras de privacidade e custo de anotação.

Em suma, o DERMAE valida que a combinação de Geração Difusiva, Pré-treinamento MAE e Destilação de Conhecimento é uma estratégia eficaz para superar a escassez de dados e o desbalanceamento de classes na dermatologia computacional.