CompDiff: Hierarchical Compositional Diffusion for Fair and Zero-Shot Intersectional Medical Image Generation

O artigo propõe o CompDiff, um framework de difusão composicional hierárquico que resolve o problema de geradores desequilibrados em imagens médicas ao decompor o condicionamento demográfico em nível de representação, permitindo a geração de imagens de alta qualidade e equitativa para subgrupos raros e interseções não vistas durante o treinamento.

O essencial

Imagine que você está tentando ensinar um artista de IA a desenhar retratos médicos (como raios-X ou imagens da retina) para ajudar a treinar médicos e diagnósticos. O problema é que, na vida real, temos muito mais fotos de pessoas brancas, jovens e do sexo masculino do que de pessoas idosas, asiáticas ou do sexo feminino.

Se você treinar o artista apenas com essas fotos desequilibradas, ele vai ficar muito bom em desenhar o que ele viu, mas vai fazer um trabalho terrível (ou nem saber desenhar) quando você pedir um retrato de uma "mulher asiática idosa".

O artigo CompDiff apresenta uma solução inteligente para esse problema. Vamos explicar como funciona usando analogias simples:

1. O Problema: O Artista Cansado e Confuso

Imagine que o modelo de IA tradicional é como um artista que recebe uma lista de pedidos escrita em um único pedaço de papel pequeno (o "prompt" de texto).

• Se você pede: "Desenhe um raio-X de uma mulher asiática de 80 anos com pneumonia", o artista tenta encaixar todas essas palavras (mulher, asiática, 80 anos) no mesmo espaço limitado.
• Como ele nunca viu uma mulher asiática de 80 anos na sua "biblioteca" de treinamento, ele fica confuso. Ele tenta adivinhar, mas o resultado sai borrado ou errado.
• Métodos antigos tentavam resolver isso dando "mais pontos" (pesos) para as fotos raras durante o treino, mas se a foto não existe no treinamento, não adianta dar mais pontos. É como tentar ensinar alguém a cozinhar um prato que nunca viu, apenas dizendo "preste mais atenção nisso".

2. A Solução: A "Caixa de Montagem" Inteligente (CompDiff)

Os autores criaram o CompDiff. A ideia principal é mudar como o artista recebe as instruções, não apenas como ele é treinado.

Eles criaram uma ferramenta chamada HCN (Rede Condicionadora Hierárquica). Pense nela como uma caixa de montagem de LEGO ou um sistema de receitas de cozinha:

• Em vez de pedir o prato pronto: O modelo não tenta memorizar "Mulher Asiática Idosa" como um bloco único e impossível.
• Ele aprende as peças separadas: Primeiro, ele aprende muito bem o que é "Idoso". Depois, o que é "Asiático". Depois, o que é "Mulher".
• Ele aprende a combinar: O modelo tem uma "receita" que diz: "Para fazer 'Idoso + Asiático', misture a peça 'Idoso' com a peça 'Asiático' de um jeito específico. Para adicionar 'Mulher', misture tudo isso com a peça 'Mulher'."

Isso permite que o modelo componha (monte) novos retratos a partir de peças que ele já conhece muito bem. Mesmo que ele nunca tenha visto uma "Mulher Asiática Idosa" antes, ele consegue montar uma imagem de alta qualidade porque entende as partes individuais e como elas se conectam.

3. Como isso funciona na prática?

• O "Token" Demográfico: O modelo cria um cartão de identidade especial para a pessoa (idade, raça, sexo) separado do texto médico.
• A Hierarquia: Ele primeiro junta idade e sexo, depois junta com a raça. É como construir uma árvore genealógica: você começa com os avós (características individuais), passa pelos pais (combinações de dois) e chega ao filho (a combinação completa).
• Resultado: O modelo consegue gerar imagens de grupos que não existiam no treinamento (Zero-Shot), como se ele tivesse uma intuição matemática sobre como as características se misturam.

4. Por que isso é importante? (Os Resultados)

Os pesquisadores testaram isso em raios-X de tórax e imagens de fundo de olho. Os resultados foram impressionantes:

• Qualidade: As imagens geradas eram mais nítidas e realistas do que as dos concorrentes.
• Justiça (Fairness): O modelo não fazia imagens ruins apenas porque a pessoa era de um grupo minoritário. A qualidade era igual para todos.
• Utilidade Real: Quando eles usaram essas imagens geradas para treinar um novo sistema de diagnóstico médico, esse sistema ficou mais preciso e menos preconceituoso. Ele não deixava de diagnosticar doenças em pessoas idosas ou de outras raças.

Resumo Final

O CompDiff é como ensinar um artista a não decorar rostos inteiros, mas sim a entender como os olhos, o nariz e a pele se combinam. Assim, quando ele precisa desenhar um rosto que nunca viu, ele não fica travado; ele pega as peças que conhece e as monta de forma criativa e precisa.

Isso é um grande passo para garantir que a Inteligência Artificial na medicina funcione bem para todos, e não apenas para a maioria das pessoas.

Resumo técnico

Resumo Técnico: CompDiff

1. O Problema: O "Problema do Gerador Desequilibrado"

O artigo identifica uma lacuna crítica no uso de modelos generativos (especificamente modelos de difusão) para síntese de imagens médicas. Embora essas ferramentas sejam promissoras para aumentar conjuntos de dados desbalanceados e treinar IA mais justa, existe uma premissa não examinada: os próprios geradores produzem imagens de qualidade igualitária entre todos os grupos demográficos?

• A Raiz do Problema: Modelos treinados em dados desbalanceados herdam essas desigualdades. Eles geram amostras de alta qualidade para grupos majoritários, mas a qualidade degrada-se drasticamente para subgrupos raros.
• Interseccionalidade Ausente: O desafio é agravado em interseções demográficas raras ou inexistentes nos dados de treinamento (ex: "paciente asiática, idosa, com uma patologia específica"). Técnicas tradicionais como reamostragem (oversampling) ou reponderação de perda (loss reweighting) falham aqui, pois não podem criar sinal de aprendizado para combinações que nunca foram observadas.
• Limitação das Soluções Atuais: Abordagens como o FairDiffusion operam no nível de otimização (reponderando a perda), mas não alteram como as demografias são representadas internamente. Elas dependem de tokens de texto no prompt, onde tokens demográficos competem pelo espaço limitado do CLIP e não conseguem generalizar para combinações não vistas.

2. Metodologia: CompDiff e a Rede Condicionadora Hierárquica (HCN)

O CompDiff propõe uma solução no nível de representação, introduzindo uma estrutura que permite a generalização composicional zero-shot. A ideia central é que a identidade demográfica é composicional: uma interseção rara pode ser construída a partir de atributos únicos bem aprendidos e interações de pares.

Arquitetura Principal:

1. Separação de Fluxos:
   • Achados clínicos são codificados via CLIP (como em modelos padrão).
   • Atributos demográficos (idade, sexo, raça) são processados separadamente por uma Rede Condicionadora Hierárquica (HCN).
2. Hierarquia da HCN:
   • Nível de Atributos Únicos ("Avós"): Cada atributo ($x_v$) é mapeado para um embedding latente ($e_v$).
   • Nível de Interações em Par ("Pais"): MLPs dedicados modelam interações não aditivas entre pares de atributos (ex: idade + sexo, idade + raça, sexo + raça). Isso captura relações complexas sem superajustar (overfitting) em subgrupos raros.
   • Nível de Composição Completa ("Filho"): Um MLP final combina as interações em pares para gerar a representação demográfica final ($h_{demo}$).
3. Integração com o Modelo de Difusão:
   • A representação latente demográfica é amostrada (via reparametrização) e projetada em um token demográfico ($c$).
   • Este token é concatenado com os embeddings de texto clínico para formar o contexto de cross-attention do UNet.
4. Objetivo de Treinamento:
   • Perda de Difusão ($L_{diff}$): Padrão.
   • Perda de Consistência Composicional ($L_{comp}$): Um "âncora" suave que garante que a representação composta não se afaste excessivamente de uma base aditiva simples, estabilizando o treinamento.
   • Perda Auxiliar ($L_{aux}$): Classificação direta aplicada no token projetado ($c$) que o UNet vê, garantindo que a informação demográfica seja preservada após a projeção.

3. Contribuições Chave

• Mudança de Paradigma: Move o foco da reponderação de dados/otimização para o design arquitetônico da representação (indução de viés estrutural).
• Generalização Zero-Shot Interseccional: Capacidade de gerar imagens de alta qualidade para combinações demográficas que não existiam no conjunto de treinamento, compondo-as a partir de atributos aprendidos.
• Eficiência de Parâmetros: A HCN adiciona apenas ~0,19% de parâmetros treináveis ao modelo base (Stable Diffusion 2.1), sem alterar os passos de difusão ou amostragem.
• Justiça sem Perda de Qualidade: Demonstra que é possível melhorar a equidade entre subgrupos sem sacrificar a qualidade geral da imagem ou o desempenho dos grupos majoritários.

4. Resultados Experimentais

Os experimentos foram realizados em duas modalidades médicas: Radiografias de Tórax (MIMIC-CXR) e Imagens de Fundo de Olho (FairGenMed).

• Qualidade de Geração:
  • O CompDiff superou tanto o Fine-tuning padrão quanto o FairDiffusion na métrica FID (Fréchet Inception Distance).
  • Exemplo (Raio-X): FID de 64.3 (CompDiff) vs. 75.1 (FairDiffusion) vs. 82.8 (Baseline).
  • Alinhamento clínico superior: Maior AUROC na classificação de doenças (0.82 vs. 0.74).
• Equidade (Fairness):
  • O CompDiff alcançou o menor ES-FID (Equity-Scaled FID), indicando menor disparidade de qualidade entre subgrupos de sexo, raça e idade.
  • Em subgrupos raros (ex: "Feminino, Asiático, 40-60 anos"), o CompDiff reduziu o FID de 204.0 (Baseline) para 167.9, enquanto o FairDiffusion teve melhoria marginal.
• Generalização Zero-Shot:
  • Ao remover 5 interseções demográficas inteiras do treinamento, o CompDiff manteve a capacidade de gerar imagens para esses grupos, com melhoria de até 21% no FID em comparação às baselines. O FairDiffusion, na verdade, piorou em alguns casos raros, confirmando que a reponderação de perda não funciona sem sinal de treinamento.
• Utilidade Descendente (Downstream Utility):
  • Classificadores de doenças treinados em dados sintéticos do CompDiff apresentaram melhor desempenho (maior AUROC) e menor viés demográfico (menor taxa de subdiagnóstico e menor diferença de Equalized Odds) ao serem avaliados em dados reais.

5. Significado e Conclusão

O artigo demonstra que o design da condicionamento demográfico é um fator crucial e subexplorado na geração de imagens médicas justas.

• Implicação Principal: A justiça na IA médica não depende apenas de quantos dados se tem, mas de como a informação demográfica é estruturada dentro do modelo. A abordagem composicional hierárquica permite que o modelo "raciocine" sobre combinações não vistas, similar à forma como modelos de linguagem compõem palavras em novas frases.
• Limitações: A avaliação ainda depende de métricas quantitativas (sem avaliação de especialistas clínicos) e a abordagem assume atributos demográficos estruturados, não se estendendo facilmente a atributos contínuos ou não estruturados.

Em suma, o CompDiff oferece uma solução arquitetônica robusta para mitigar o viés interseccional na geração de imagens médicas, permitindo a criação de conjuntos de dados sintéticos mais justos e úteis para o treinamento de sistemas de diagnóstico robustos.