Steering Diffusion Models via Class-Contrastive Influence for Few-Shot Medical Classification

Resumo Técnico: Direcionando Modelos de Difusão via Influência Contrastiva de Classe para Classificação Médica de Poucos Disparos (Few-Shot)

Definição do Problema
Na classificação de imagens médicas com poucos disparos (few-shot), os dados rotulados são escassos, tornando o aumento de dados essencial. Embora modelos de difusão prontos para uso possam gerar amostras de treinamento sintéticas, as abordagens existentes concentram-se principalmente em melhorar o realismo, a diversidade ou a adaptação de domínio dessas imagens. Os autores argumentam que essas métricas não garantem a utilidade para a tarefa de classificação subsequente. Existe uma discrepância significativa onde imagens geradas com qualidade visual comparável produzem desempenhos de classificação vastamente diferentes. O problema central identificado é a falta de um critério fundamentado para medir e otimizar a "utilidade da tarefa" das amostras sintéticas. O artigo postula que a limitação central das práticas atuais não é apenas a qualidade da geração, mas a falha em gerar amostras que ajudem especificamente o classificador a refinar sua fronteira de decisão.

Metodologia
A solução proposta, Class-Contrastive Influence (C2I - Influência Contrastiva de Classe), desloca o objetivo de otimização da qualidade da imagem para a influência baseada em gradiente das amostras sintéticas no classificador.

1. Critério de Influência Contrastiva de Classe (C2I):

   • O método utiliza a estimativa de influência baseada em gradiente para quantificar como uma amostra de treinamento afeta a perda de validação.
   • Os autores observam que pontuações de influência brutas e agnósticas à classe não se correlacionam com o desempenho da classificação porque os conjuntos de validação contêm sinais conflitantes de diferentes classes.
   • Em vez disso, eles definem o C2I com base no hiato de influência (influence gap): uma amostra é considerada útil se seu gradiente de perda se alinha positivamente com os gradientes de validação da mesma classe e se opõe (alinha-se negativamente com) aos gradientes de outras classes.
   • Para classificação binária, o C2I é formulado como a diferença ao quadrado entre a média da influência nos conjuntos de validação positivo e negativo, normalizada pela variância. Para tarefas multiclasse, é definido como o escore softmax da média da influência na classe correspondente em relação a todas as classes.

2. Conexão Teórica com a Dificuldade da Amostra (Sample Hardness):

   • A análise teórica (Teorema 1) demonstra que maximizar o C2I em um cenário de regressão logística conduz a amostra gerada em direção à média global do conjunto de validação.
   • Como a média global de dois agrupamentos de classes normalmente reside próxima à fronteira de decisão, maximizar o C2I inerentemente seleciona exemplos "difíceis" — amostras que são difíceis de classificar e residem próximas à fronteira.
   • Treinar nessas amostras difíceis e próximas à fronteira ajuda a refinar a fronteira de decisão, melhorando a generalização e a robustez.

3. Ajuste Fino por Aprendizado por Reforço (RL):

   • Para operacionalizar isso, os autores propõem um esquema de ajuste fino para modelos de difusão de Texto-para-Imagem (T2I) pré-treinados (especificamente o Stable Diffusion 2.1).
   • Função de Recompensa: O escore C2I serve como o sinal de recompensa em um framework de Aprendizado por Reforço. O gerador é ajustado para maximizar a recompensa C2I esperada.
   • Otimização: Os autores utilizam a Otimização de Política de Difusão de Denoising (DDPO) para atualizar os parâmetros do modelo de difusão (especificamente os pesos LoRA) com base na recompensa C2I. Isso direciona o processo de geração para as regiões "informativas de classe" do manifold de dados sem exigir mudanças arquiteturais no gerador ou no classificador.
   • Fluxo de Trabalho: Um classificador (ViT-B/16) é primeiro ajustado no conjunto de treinamento de poucos disparos. Seus gradientes em um conjunto de validação mantido à parte são usados para calcular a recompensa C2I para lotes gerados, que então guiam o ajuste fino por RL do modelo de difusão.

Principais Contribuições

• Formalização da Eficácia: O artigo formaliza a eficácia da amostra através do C2I, demonstrando que amostras de alto C2I correspondem a exemplos difíceis e próximos à fronteira, em vez de apenas imagens realistas.
• Esquema de Geração Baseado em RL: Introduz um framework de RL plug-and-play que utiliza o C2I como recompensa para direcionar modelos de difusão prontos para o uso visando gerar amostras úteis para a tarefa e informativas de classe.
• Validação Empírica: O método é validado em vários benchmarks de imagens médicas de poucos disparos (BreastMNIST, DermaMNIST, PneumoniaMNIST), mostrando melhorias consistentes sobre o aumento de dados padrão e outros baselines baseados em difusão.

Resultados Experimentais
O método foi avaliado em três benchmarks MedMNIST sob configurações de poucos disparos (16 ou 32 amostras rotuladas por classe).

• Desempenho: O gerador guiado por C2I superou consistentemente os baselines, incluindo transformações padrão (RandAugment, RandomErasing, Mixup) e outros métodos de aumento baseados em difusão (DataDream, Dataset Expansion, DistDiff).
  • Em média, o método proposto alcançou o maior AUC entre diferentes backbones (ViT-B/16 e ResNet-18).
  • Notavelmente, o método melhorou o desempenho no ResNet-18, indicando que as amostras sintetizadas são amplamente informativas e transferem entre arquiteturas de modelo.
• Robustez: Modelos treinados com aumento guiado por C2I demonstraram robustez superior sob várias condições de ruído (sal e pimenta, desfoque gaussiano, compressão JPEG), sugerindo a geração de uma fronteira de decisão mais estável.
• Correlação de Dificuldade: Experimentos confirmaram que, conforme o ajuste fino por RL progride, a recompensa C2I aumenta e as características das imagens geradas movem-se para mais perto dos centroides de características de validação, confirmando a geração de exemplos mais difíceis.
• Ablação: Aumentar o número de épocas de treinamento por RL e o número de imagens sintetizadas geralmente melhorou o AUC subsequente, enquanto alguns baselines (por exemplo, Dataset Expansion) mostraram degradação de desempenho com o aumento de dados sintéticos.

Significância e Alegações
O artigo alega que o aumento sintético é mais eficaz quando guiado pela utilidade da tarefa em vez de apenas pela qualidade da imagem. Ao alinhar os dados sintéticos com os gradientes de validação, o C2I transforma modelos de difusão genéricos em geradores direcionados que fortalecem as fronteiras de decisão onde os dados rotulados são mais escassos.

Os autores reconhecem limitações, observando que seu método introduz um custo computacional adicional em comparação com o uso apenas dos dados de treinamento originais, principalmente devido ao processo de ajuste fino por RL (exigindo aproximadamente 5 horas de GPU em uma A100). No entanto, eles afirmam que esse custo é justificado pelos ganhos consistentes em precisão e robustez em regimes de baixos dados. O estudo sugere uma mudança de paradigma no uso de modelos de difusão para aumento: otimizar a geração para sua influência subsequente no classificador.