DMin: Scalable Training Data Influence Estimation for Diffusion Models

1. O Problema

Os modelos de difusão (Diffusion Models - DMs) são ferramentas poderosas para geração de imagens, mas são treinados em conjuntos de dados massivos e diversos da internet. Isso levanta questões críticas sobre transparência e viés: é difícil entender quais amostras específicas do conjunto de treinamento influenciaram uma imagem gerada.

Existem métodos existentes para estimar essa influência (como DataInf, K-FAC, TRAK), mas eles enfrentam barreiras intransponíveis quando aplicados a modelos de grande escala (bilhões de parâmetros) e grandes conjuntos de dados:

• Limitação de Armazenamento: Métodos de segunda ordem (que exigem inversão da matriz Hessiana) ou métodos de primeira ordem com projeção aleatória exigem armazenar gradientes de todos os dados de treinamento. Para um modelo como o Stable Diffusion 3 Medium (2 bilhões de parâmetros), armazenar gradientes brutos para 10.000 amostras em 10 passos de tempo exigiria cerca de 800 TB de armazenamento, o que é impraticável.
• Limitação Computacional: A necessidade de calcular e armazenar matrizes de projeção gigantes ou realizar inversões de Hessiana torna esses métodos lentos e inviáveis para modelos grandes ou modelos com LoRA (Low-Rank Adaptation) em escala.
• Fragilidade: Estimações de influência em modelos profundos tendem a ser instáveis devido a gradientes com magnitudes excessivamente grandes.

2. Metodologia: DMin

O artigo propõe o DMin (Diffusion Model influence), um framework escalável para estimar a influência de dados de treinamento em imagens geradas, capaz de lidar com modelos de bilhões de parâmetros. A abordagem baseia-se em três pilares principais:

A. Cálculo e Coleta de Gradientes

Para cada amostra de treinamento $(x_i, p_i)$, o DMin calcula os gradientes do modelo em relação à função de perda de previsão de ruído em vários passos de tempo ($t$).

• Para modelos com adaptadores (como LoRA), apenas os gradientes dos parâmetros do adaptador são coletados, reduzindo drasticamente o volume de dados.
• Para modelos completos, todos os gradientes são coletados.

B. Compressão de Gradientes (O Núcleo da Escalabilidade)

Para resolver o problema de armazenamento, o DMin aplica uma técnica de compressão de quatro etapas nos vetores de gradiente, reduzindo o tamanho de gigabytes para kilobytes:

1. Padding: O vetor de gradiente é preenchido com zeros até um tamanho divisível por um fator $v$.
2. Permutação: Uma permutação aleatória é aplicada para quebrar qualquer estrutura inerente no vetor.
3. Projeção Aleatória (Random Projection): O vetor é multiplicado elemento a elemento por um vetor de projeção binário (elementos aleatórios de -1 ou 1).
4. Agrupamento (Group Addition): Os elementos do vetor são divididos em grupos e somados, produzindo um vetor comprimido de dimensão $v$.

Resultado da Compressão: Em vez de armazenar o vetor de gradiente completo (ex: 2 bilhões de parâmetros), o DMin armazena apenas o vetor de permutação e o vetor de projeção binária. Isso reduz o requisito de armazenamento de centenas de Terabytes para Megabytes ou Kilobytes por amostra (ex: de ~40 GB para ~80 KB).

Normalização: Para mitigar a instabilidade de gradientes em modelos profundos, os vetores são normalizados via norma L2 antes da compressão, evitando que valores extremos dominem o produto interno.

C. Estimativa de Influência e Busca KNN

• Estimativa Exata: A influência é calculada como o produto interno entre os gradientes comprimidos da imagem gerada e os gradientes comprimidos cacheados das amostras de treinamento.
• Busca KNN (K-Nearest Neighbors): Para cenários onde apenas as $k$ amostras mais influentes são necessárias, o DMin constrói um índice KNN (usando o algoritmo HNSW) sobre os gradientes comprimidos. Isso permite recuperar as amostras mais relevantes em menos de 1 segundo, sem precisar calcular produtos internos com todas as amostras.

3. Contribuições Principais

1. Primeiro Método Escalável para Modelos Grandes: É a primeira técnica capaz de realizar estimativa de influência em modelos de difusão com bilhões de parâmetros (como SD 3 Medium Full).
2. Redução Extrema de Armazenamento: A técnica de compressão permite armazenar gradientes de todo um conjunto de treinamento em uma fração mínima do espaço original (de TBs para MBs/KBs), mantendo a precisão.
3. Velocidade de Inferência: A combinação de compressão e busca KNN permite recuperar as top-$k$ amostras influentes em frações de segundo.
4. Implementação Open Source: O código PyTorch com suporte a multiprocessamento foi disponibilizado.

4. Resultados Experimentais

Os autores avaliaram o DMin em modelos condicionais (SD 1.4 LoRA, SD 3 Medium LoRA, SD 3 Medium Full) e incondicionais (DDPM no MNIST e CIFAR-10).

• Precisão (Taxa de Detecção): Em tarefas onde o modelo gera imagens baseadas em prompts específicos (ex: "cartas mágicas"), o DMin identificou corretamente que as amostras de treinamento mais influentes vinham do subconjunto correspondente ("Magic Cards").
  • No SD 3 Medium (Full), o DMin alcançou taxas de detecção superiores a 96% para o Top-5 em conjuntos de dados específicos, superando significativamente baselines como SSIM, CLIP, LiSSA e D-TRAK (que não puderam ser executados em modelos completos devido a limitações de memória).
  • A normalização e a compressão foram cruciais; sem normalização, a performance caiu drasticamente devido à instabilidade dos gradientes.
• Eficiência de Tempo:
  • Para o SD 3 Medium (Full), métodos anteriores levariam horas ou seriam impossíveis. O DMin com KNN (Top-5) levou apenas 0.009 segundos por amostra de teste.
  • Comparado ao método sem compressão, o DMin com KNN oferece um speedup de mais de 100.000x.
• Eficiência de Memória:
  • Para o SD 3 Medium (Full), o armazenamento necessário caiu de 339 TB (sem compressão) para apenas 726 MB (com compressão $v=2^{12}$), uma redução de 0,00017%.

5. Significado e Impacto

O DMin representa um avanço fundamental na interpretabilidade de modelos generativos de grande escala.

• Transparência e Segurança: Permite rastrear a origem de conteúdos gerados, identificando se uma imagem gerada foi influenciada por dados enviesados, prejudiciais ou protegidos por direitos autorais.
• Viabilidade Prática: Remove a barreira de hardware que impedia a análise de influência em modelos modernos (bilhões de parâmetros), tornando-a acessível em hardware de consumo ou clusters padrão.
• Futuro da Pesquisa: Estabelece um novo padrão para métodos de influência em modelos profundos, demonstrando que técnicas de compressão inteligente podem superar limitações de armazenamento sem sacrificar a precisão da análise.

Em resumo, o DMin torna possível "abrir a caixa preta" dos maiores modelos de difusão atuais, permitindo que pesquisadores e desenvolvedores entendam e auditem como esses modelos aprendem e geram conteúdo.