Geometry-Aware Dataset Condensation for Diffusion Model Training

Resumo Técnico: Condensação de Conjuntos de Dados Sensível à Geometria para Treinamento de Modelos de Difusão

1. Declaração do Problema

A condensação de conjuntos de dados visa construir conjuntos de dados compactos a partir de dados reais via síntese ou seleção para reduzir os custos de armazenamento e computação associados ao treinamento de modelos de aprendizado profundo. Embora os métodos de condensação existentes tenham mostrado sucesso em tarefas discriminativas (ex: classificação), eles são inadequados para o treinamento de modelos de difusão.

• Limitações de Síntese: Métodos que geram dados sintéticos via otimização de pixels contínua frequentemente produzem amostras de baixa fidelidade. Os modelos de difusão são altamente sensíveis ao ruído e distorções estruturais, tornando essas amostras sintéticas inadequadas para uma modelagem autêntica.
• Limitações de Seleção: Os métodos existentes de seleção de subconjuntos reais geralmente dependem de critérios fixos ou classificações heurísticas (ex: o D2C classifica imagens por dificuldade de difusão ao longo de um eixo unidimensional). Essas abordagens falham em preservar a geometria de distribuição complexa exigida pelos objetivos de verossimilhança da difusão. Eles frequentemente carecem de um objetivo de otimização fundamentado e alinhado com o treinamento de difusão, resultando em subconjuntos que são mal alinhados com a distribuição total dos dados e falham em capturar a estrutura de variedade necessária.

O desafio central é selecionar um subconjunto real compacto que preserve a estrutura geométrica e a fidelidade distributiva do conjunto de dados completo, especificamente otimizado para os objetivos de treinamento baseados em verossimilhança de modelos de difusão.

2. Metodologia

Os autores propõem a Condensação de Conjunto de Dados Sensível à Geometria (GADC), que reformula a seleção de subconjuntos reais como um problema de alinhamento de distribuição sensível à geometria. O framework consiste em três componentes principais:

A. Transporte Ótimo Parcial Unilateral (POT)
Para abordar a incompatibilidade de capacidade entre um pequeno subconjunto e o conjunto de dados completo, os autores empregam o transporte ótimo parcial unilateral.

• Conceito: Ao contrário do OT equilibrado que impõe o alinhamento de toda a massa, o POT unilateral permite que a massa de origem (subconjunto selecionado) seja totalmente transportada enquanto relaxa a massa de destino (conjunto de dados completo) sob uma restrição de capacidade.
• Mecanismo: Isso permite massa não correspondente em regiões periféricas de baixa densidade, concentrando o transporte em regiões dominantes, geometricamente estáveis e de alta densidade do manifold de dados.
• Implementação: O problema é resolvido eficientemente usando uma reformulação de "fonte-dummy" e regularização entrópica, otimizada via iterações de Sinkhorn. Isso permite uma alocação de massa flexível onde o subconjunto foca em regiões representativas enquanto mantém a cobertura necessária.

B. Regularização Estatística e Semântica
Para complementar o alinhamento geométrico e garantir a fidelidade distributiva, dois termos de regularização leves são introduzidos:

• Regularização de Média-Variância ($L_{sta}$): Alinha as estatísticas de primeira e segunda ordem (média e variância) das representações de características entre o subconjunto selecionado e o conjunto de dados completo para preservar a fidelidade distributiva global.
• Regularização de Confiança ($L_{conf}$): Garante a consistência semântica ao penalizar a seleção de amostras com baixas probabilidades de classe previstas. Isso evita a inclusão de âncoras geométricas não confiáveis que poderiam comprometer o alinhamento.

C. Estratégia de Otimização Discreta em Dois Estágios
Como a seleção de um subconjunto de tamanho fixo é um problema combinatório, os autores propõem um solver eficiente de dois estágios:

1. Estágio I (Seleção Gananciosa Guiada pela Geometria): Um processo de construção incremental que seleciona amostras para minimizar o objetivo composto (perda POT + regularizações), estabelecendo uma ampla cobertura do manifold.
2. Estágio II (Refinamento Baseado em Trocas): Uma fase de refinamento que realiza trocas par a par entre amostras selecionadas e não selecionadas para corrigir escolhas miopes iniciais, otimizando ainda mais o alinhamento geométrico global e a consistência.

3. Principais Contribuições

1. Reformulação do Problema: O artigo reformula a condensação de conjuntos de dados para modelos de difusão como um problema de alinhamento de distribuição, afastando-se de classificações heurísticas ou baseadas em escalares. Introduz um objetivo de POT unilateral com regularização estatística para selecionar subconjuntos reais representativos.
2. Framework de Otimização: Propõe um framework de otimização discreta de dois estágios (construção gananciosa seguida de refinamento baseado em trocas) que resolve eficientemente o problema de alinhamento, evitando as limitações de amostragem baseada em critérios fixos ou de classificação.
3. Desempenho Empírico: Experimentos extensos demonstram que o método alcança fidelidade e cobertura distributiva superiores em comparação com abordagens anteriores em várias variantes de difusão (DiT, SiT), tamanhos de subconjunto, resoluções de imagem e rodadas de treinamento.

4. Resultados Experimentais

O método foi avaliado no ImageNet-1K usando arquiteturas DiT-L/2 e SiT-L/2 com orçamentos de dados variando de 0,8% (10K imagens) a 8% (100K imagens).

• Métricas de Desempenho: O método proposto alcançou consistentemente a menor Distância de Fréchet Inception (FID) e os maiores Inception Score (IS), Precisão e Recall em comparação com os baselines, incluindo Random Sampling, K-Center, Herding, CCS, DQ e D2C.
  • Exemplo: No ImageNet 256×256 com um orçamento de 10K dados, o método alcançou um FID de 3,43, superando o D2C (4,20) e o Random (35,86).
  • Exemplo: Na resolução 512×512 com um orçamento de 10K, o método alcançou um FID de 6,17, significativamente melhor que o D2C (14,8).
• Robustez: As melhorias persistiram através de diferentes variantes de difusão (DiT e SiT), diferentes protocolos de avaliação (10K vs. 50K amostras geradas) e extensas iterações de treinamento (até 300K).
• Eficiência: O processo de seleção é computacionalmente eficiente. Em uma única GPU, o método levou 5,5 horas para selecionar um subconjunto de 10K do ImageNet, comparado a 41,9 horas para o D2C e 30,4 horas para o DQ.

Estudos de Ablação:

• A remoção da perda POT ($L_{OT}$) degradou o FID e reduziu a precisão/recall, confirmando a necessidade do alinhamento geométrico.
• Substituir o POT unilateral pelo OT equilibrado degradou o desempenho, validando o benefício de permitir massa não correspondente em regiões periféricas.
• A remoção da regularização estatística ($L_{sta}$) ou de confiança ($L_{conf}$) levou a um pior FID ou IS, respectivamente, destacando a importância das estatísticas globais e da clareza semântica.

5. Significância e Alegações

O artigo afirma que, ao tratar a condensação de conjuntos de dados como um problema de alinhamento de distribuição sensível à geometria, ele consegue preencher a lacuna entre a seleção de dados e os requisitos específicos do treinamento de modelos de difusão.

• Eficiência de Dados: O método permite modelagem generativa de alta qualidade em conjuntos de dados condensados, reduzindo significativamente os custos de armazenamento e computação associados ao treinamento de modelos de difusão.
• Fidelidade: Ao contrário dos métodos de geração sintética, a abordagem retém a estrutura de alta fidelidade das amostras reais enquanto garante que o subconjunto se alinhe com a geometria da distribuição total dos dados.
• Praticidade: A estratégia de otimização de dois estágios proposta é escalável e computacionalmente eficiente, tornando-a uma solução prática para cenários de recursos limitados.

Os autores concluem que seu método fornece um substrato de dados confiável para o treinamento generativo do tipo difusão, oferecendo um caminho complementar e centrado em dados para melhorar a eficiência junto com as otimizações do lado do modelo.