An exact information theory of generalization phase transitions in Bayesian diffusion models

Título: Uma Teoria Exata da Teoria da Informação sobre Transições de Fase de Generalização em Modelos de Difusão Bayesiana

Enunciado do Problema
Um mistério fundamental na IA generativa é como os modelos de difusão contornam a "maldição da dimensionalidade" para aprender distribuições complexas sobre espaços de alta dimensão a partir de conjuntos de dados finitos sem meramente memorizá-los. Naivamente, a transição da memorização para a generalização para distribuições de dados arbitrárias deveria exigir uma quantidade de dados intratável (exponencial na dimensão). Embora modelos de score empíricos (que revertem a difusão usando o score empírico exato de um conjunto de dados finito) sejam conhecidos por memorizar, modelos de difusão neurais treinados frequentemente generalizam de forma robusta. Abordagens teóricas existentes têm tido dificuldade em fornecer uma teoria precisa e preditiva de quando e onde a transição da memorização para a generalização ocorre para modelos de difusão não lineares treinados através de diversas arquiteturas e conjuntos de dados.

Metodologia
Os autores introduzem uma nova classe de modelos analiticamente tratáveis chamados Modelos de Difusão com Restrição de Informação Bayesiana (BIRD - Bayesian Information Restricted Diffusion).

• Conceito Central: Nos modelos BIRD, cada pixel $x$ atua como um agente que observa apenas uma parte restrita, possivelmente estocástica, da imagem ruidosa $\phi_t$ (denotada por $C_{x,t}$), em vez da imagem completa. O pixel então utiliza inferência Bayesiana ótima para adivinhar qual amostra de treinamento passada $\phi \in \mathcal{D}$ produziu essa observação restrita.
• Mecanismo: O modelo computa um valor de pixel desruído como a média posterior sobre o conjunto de treinamento, ponderada pela probabilidade posterior $P_{\text{train}}(\phi | C_{x,t})$. Ao restringir a informação disponível para cada pixel (por exemplo, observando apenas um patch espacial local), o "jogo de adivinhação Bayesiana" torna-se mais difícil, teoricamente impedindo o modelo de colapsar para um único ponto de treinamento (memorização).
• Estrutura Teórica: Os autores analisam a transição memorização-generalização usando a teoria da informação. Eles derivam uma fronteira de fase no espaço conjunto de:
  1. A quantidade de dados de treinamento ($|\mathcal{D}|$).
  2. O tempo/nível de ruído no processo reverso ($t$ ou $\sigma_t$).
  3. A quantidade de restrição de informação (por exemplo, o tamanho do patch $L$).
• Critério Chave: A transição ocorre quando a informação mútua $I(\phi; C_{x,t})$ entre a observação restrita e a verdadeira distribuição de dados é igual ao logaritmo do tamanho do conjunto de dados: $\ln |\mathcal{D}| = I(\phi; C_{x,t})$.
  • Fase de Memorização: $I(\phi; C_{x,t}) > \ln |\mathcal{D}|$ (A observação é informativa demais; o modelo adivinha o ponto de treinamento perfeitamente).
  • Fase de Generalização: $I(\phi; C_{x,t}) < \ln |\mathcal{D}|$ (A observação é insuficiente para identificar unicamente um ponto de treinamento; o modelo faz a média sobre as possibilidades).

Principais Contribuições

1. Introdução dos Modelos BIRD: Um framework geral e analiticamente solucionável que engloba e generaliza modelos de score locais (LS), modelos de score equivariantes (ES) e modelos de score local-equivariantes (ELS). Ele não requer suposições explícitas de arquitetura ou dinâmica de aprendizado, baseando-se puramente na inferência Bayesiana sobre um conjunto de dados finito com restrição de informação.
2. Predição da Dinâmica Inicial de Treinamento: Os autores demonstram que modelos BIRD espacialmente locais (onde $C_{x,t}$ é um patch de imagem local) aproximam-se de perto o comportamento de modelos de difusão neurais treinados (incluindo UNets e Diffusion Transformers) durante os estágios iniciais de treinamento. Eles alcançam alta precisão preditiva (mediana de $r^2 \sim 0.85-0.93$) em saídas de imagens individuais através de várias arquiteturas e conjuntos de dados (Celeba64, CIFAR10, MNIST, FashionMNIST).
3. Teoria da Fronteira de Fase Exata: Sob suposições mínimas sobre as distribuições de dados, o artigo deriva um critério informacional preciso para a transição memorização-generalização. Eles provam que a generalização ocorre se, e somente se, a informação mútua entre a observação restrita e os dados for menor que o logaritmo do número de pontos de treinamento.
4. Escala de Patch Crítica e Alinhamento Espectral: Para modelos espacialmente locais, os autores identificam uma escala de patch crítica dependente do tempo $L_c(t)$. Eles mostram que a geração ótima procede próximo à borda da memorização, onde a escala do patch $L$ acompanha a escala crítica $L_c$. Eles também demonstram que, para imagens naturais com densidades espectrais de lei de potência ($P(k) \sim k^{-2-\epsilon}$), a escala crítica alinha-se com a "escala espectral" (a escala de denoising linear ótimo).
5. Contorno da Maldição da Dimensionalidade: A teoria revela que, para imagens invariantes de escala ($\epsilon=0$), os modelos BIRD não sofrem com a maldição da dimensionalidade; o tamanho de dados necessário não precisa escalar com o tamanho da imagem. Para imagens quase invariantes de escala ($\epsilon \sim 0.1-0.3$), o requisito de dados escala como $\exp(L_I^\epsilon)$, o que é significativamente sub-exponencial comparado ao requisito $\exp(L_I^2)$ ingênuo.

Resultos

• Validação Empírica: Experimentos confirmam a localização da transição de fase prevista teoricamente. As curvas de informação mútua derivadas da teoria seguem de perto o déficit de entropia observado em experimentos numéricos tanto em conjuntos de dados de brinquedo quanto em conjuntos de dados de imagens do mundo real.
• Consistência: Tanto os modelos BIRD quanto as redes neurais treinadas precocemente exibem "generalização consistente", onde modelos treinados em subconjuntos distintos de dados produzem saídas quase idênticas quando recebem o mesmo input de ruído.
• Leis de Escala: A escala de patch crítica $L_c(\sigma_t)$ mostra-se diminuir monotonicamente conforme o ruído diminui (o tempo progride no processo reverso), alinhando-se com o processo de geração de grosso-para-fino observado. A escala crítica é encontrada altamente correlacionada com a escala espectral $L_{\text{spec}}$ derivada da densidade espectral de potência dos dados.
• Modos de Falha: O artigo ilustra que, se a escala do patch for muito pequena ($L \ll L_c$), o modelo generaliza mas performa mal (alto erro de denoising). Se a escala do patch for muito grande ($L > L_c$), o modelo memoriza os dados de treinamento e falha em generalizar para dados de teste.

Significância
O artigo afirma fornecer um insight fundamental sobre como a IA generativa contorna a maldição da dimensionalidade. Ao enquadrar o problema através da lente da restrição de informação, os autores mostram que a generalização bem-sucedida em modelos de difusão não é meramente um resultado de vieses indutivos arquiteturais (como suavidade ou baixa dimensionalidade), mas é fundamentalmente impulsionada pela limitação de informação disponível para o processo generativo em qualquer etapa. Os resultados sugerem que a geração procede "na borda da memorização", onde a restrição de informação é cuidadosamente ajustada para evitar que o modelo colapse para exemplos específicos de treinamento enquanto retém informação suficiente para reconstruir a distribuição de dados subjacente. Este trabalho oferece um novo laboratório teórico para estudar a transição memorização-generalização através de uma ampla variedade de distribuições de dados e classes de modelos.