Flow Matching with Optimized Subclass Priors for Medical Image Augmentation

Resumo Técnico: Flow Matching com Priors de Subclasses Otimizadas para Aumento de Imagens Médicas

Declaração do Problema
Conjuntos de dados de imagem médica são caracterizados por desequilíbrios severos de cauda longa, onde condições comuns dominam enquanto patologias raras e clinicamente críticas (por exemplo, pneumotórax pequenos, bronquiectasia) são subrepresentadas. Esse desequilíbrio faz com que classificadores discriminativos superajustem-se às classes principais e falhem em condições raras. Embora o aumento de dados generativo ofereça uma solução potencial, o flow matching (FM) condicional padrão e os modelos de difusão enfrentam duas dificuldades cumulativas quando treinados em rótulos de doenças grosseiros:

1. Agregação Multimodal: Um único rótulo diagnóstico frequentemente agrega subtipos diversos, protocolos de aquisição e comorbidades. Isso cria condicionais multimodais que enviesam os geradores em direção a submodos dominantes, efetivamente truncando variações raras, mas clinicamente válidas.
2. Geometria de Transporte Ineficiente: Traçar todas as trajetórias de geração a partir de uma fonte Gaussiana padrão compartilhada força subpopulações específicas através de caminhos de transporte desproporcionalmente longos. Isso infla a dificuldade de regressão e degrada a fidelidade da geração de classes de cauda.

Metodologia
Os autores propõem uma estratégia offline que introduz priors informativos em dois níveis para refinar o processo generativo sem modificar o loop de treinamento central do modelo de flow matching.

1. Indução de Subclasses via Modelagem de Mistura Gaussiana (GMM):

   • O método particiona cada rótulo grosseiro $c$ em submodos coerentes (subclasses $k$) dentro do espaço latente do modelo generativo.
   • Utilizando um VAE congelado (FLUX1.dev), latentes determinísticos são computados. Para cada classe $c$, um Modelo de Mistura Gaussiana diagonal é ajustado aos latentes residuais (após subtrair o centro da classe $\mu_c$).
   • O número de componentes $K_c$ é selecionado usando o Critério de Informação Bayesiano Estendido (EBIC). As amostras recebem atribuições rígidas às subclasses, criando uma partição com variância condicional menor.
   • Teoricamente, os autores provam que condicionar-se a esses rótulos refinados não pode aumentar o risco de flow matching ótimo de Bayes; a melhoria equivale à variabilidade esperada entre subclasses da média de deslocamento condicional.

2. Otimização de Fonte Dependente de Subclasse:

   • Em vez de uma fonte Gaussiana compartilhada, o método aprende uma distribuição de fonte específica $\pi_0(\cdot | c, k)$ para cada subclass. Essas fontes são parametrizadas como Gaussianas aprendíveis $N(\mu_{c,k}, \text{diag}(\sigma^2_{c,k}))$ que realinham e reescalam a distribuição inicial.
   • Isso encurta as trajetórias de transporte típicas e reduz a dispersão dentro da subclass.
   • Controle Geométrico: Para prevenir soluções degeneradas (como inflação de norma ou caminhos excessivamente concentrados), os autores otimizam um objetivo proxy offline que impõe:
     • Compactação Direcional: Incentivar que as direções de deslocamento normalizadas se alinhem moderadamente com protótipos aprendíveis.
     • Limitação do Comprimento do Caminho: Impor uma penalidade fixa em outliers usando um limite específico de subclass ($r_{cap}$), definido como o percentil 99 das distâncias ao centro da classe. Isso impede que o modelo contorne a perda inflando normas de deslocamento.
     • Regularização de Escala: Regularizar as escalas logarítmicas das distribuições de fonte em direção a uma referência unitária.

Principais Contribuições

1. Formulação Condicional por Subclasse: Uma abordagem inovadora para geração de imagens médicas de cauda longa que particiona condicionais multimodais via atribuições rígidas induzidas por GMM, efetivamente reduzindo a ambiguidade condicional.
2. Fontes Dependentes de Subclasse Aprendidas: A introdução de priors informativos que remodelam o transporte condicional ao realinhar e reescalar a distribuição de fonte por subclass, reduzindo assim a dispersão dentro da subclass.
3. Objetivo Regularizado Geometricamente: Um objetivo proxy prático que equilibra compactação angular moderada com controle explícito de outliers de comprimento de caminho, prevenindo atalhos de inflação de norma e colapso direcional.
4. Validação Empírica: Demonstração de melhoria na cobertura de cauda e no desempenho de classificação de cauda longa a jusante em benchmarks de radiografia torácica (MIMIC-LT, NIH-LT) e TC (CT-RATE).

Resultados Experimentais
O método foi avaliado em três benchmarks de imagem médica de cauda longa:

• Classificação a Jusante: O método proposto alcançou consistentemente as maiores taxas de precisão balanceada (bAcc) e pontuações macro-F1 em todos os três conjuntos de dados, comparado a baselines sem aumento e outras estratégias de aumento (Vanilla FM, CBFM, CPD). Notavelmente, foi a única estratégia a melhorar confiavelmente o desempenho em ambas as modalidades de Raio-X e TC.
• Qualidade Generativa: No MIMIC-LT, o método alcançou a melhor Distância Fréchet Inception (FID) e Pontuação de Razão Inception (IRS), indicando ganhos simultâneos em fidelidade e diversidade. No CT-RATE, embora o FID tenha sido ligeiramente superior ao de uma baseline, o método manteve alto IRS e classificação a jusante superior, sugerindo melhor cobertura distribucional em vez de mera memorização.
• Estudos de Ablação: Os resultados confirmaram que ambas as etapas do pipeline (subclasse e otimização de fonte) são necessárias. A subclasse isolada melhorou o desempenho, mas a combinação produziu os melhores resultados. A otimização sem subclasse às vezes degradou o desempenho ao fazer média sobre modos conflitantes.
• Validade da Subclasse: A análise mostrou que as subclasses induzidas capturam estrutura relevante para a patologia, e não apenas metadados de aquisição. Probes de confundidores indicaram que os metadados eram detectáveis, mas não dominantes, e a pureza kNN das subclasses aprendidas excedeu significativamente partições aleatórias.

Significado e Alegações
O artigo alega que o design de fonte e o refinamento condicional são alavancas eficazes para aumento generativo robusto em cenários de cauda longa clínica. Ao particionar rótulos heterogêneos e aprender fontes dependentes de subclass, o método reduz a ambiguidade condicional e estabiliza as trajetórias de transporte a um custo computacional negligenciável (todo o processo é offline e leva menos de uma hora de GPU por conjunto de dados). Os autores posicionam isso como uma estratégia principiante que melhora a fidelidade e a diversidade da geração de classes de cauda, levando a classificadores a jusante mais confiáveis para condições médicas raras. O trabalho destaca que a geração condicional ingênua frequentemente falha em subpopulações raras devido a problemas de geometria de transporte, que podem ser mitigados através de priors estruturados.