Latent Stochastic Interpolants

Título: Latent Stochastic Interpolants (Interpolantes Estocásticos Latentes)

1. O Problema

Os modelos de difusão e os Interpolantes Estocásticos (SI) (Stochastic Interpolants) alcançaram sucesso notável na modelagem de distribuições de dados complexas. No entanto, a aplicação direta de SI enfrenta duas limitações principais quando se trata de modelos de variáveis latentes:

1. Dependência de Espaços de Observação: O framework SI original assume que as amostras das distribuições inicial ($p_0$) e alvo ($p_1$) são diretamente observáveis no espaço de dados (ex: pixels de imagem). Isso impede o uso em modelos onde o processo generativo ocorre em um espaço latente não observado, que é aprendido conjuntamente com um codificador e um decodificador.
2. Ineficiência Computacional e Priors Simples: Aplicar SI diretamente em espaços de observação de alta dimensão (como imagens de alta resolução) é computacionalmente custoso. Além disso, muitos modelos latentes tradicionais (como VAEs) são limitados pelo uso de priors simples (ex: Gaussiana isotrópica), o que restringe a expressividade do modelo generativo.

A necessidade é de um framework que permita o aprendizado conjunto (end-to-end) de um codificador, um decodificador e um modelo generativo no espaço latente, mantendo a flexibilidade de transformar distribuições arbitrárias (como no SI) e a eficiência de operar em dimensões reduzidas.

2. Metodologia

Os autores propõem os Latent Stochastic Interpolants (LSI), um framework que estende os SI para o espaço latente através de uma otimização conjunta.

A. Formulação do Modelo

O modelo consiste em três componentes treinados simultaneamente:

1. Codificador ($p_\theta(z_1|x_1)$): Mapeia observações de alta dimensão ($x$) para um espaço latente ($z$).
2. Decodificador ($p_\theta(x_1|z_1)$): Reconstrói a observação a partir da representação latente.
3. Modelo SI Latente: Opera inteiramente no espaço latente, transformando uma distribuição prior ($p_0(z_0)$) na posterior agregada definida pelo codificador ($p_1(z_1)$).

B. Derivação do ELBO em Tempo Contínuo

A inovação central é a derivação de uma Limitante Inferior da Evidência (ELBO) principial diretamente em tempo contínuo para variáveis latentes dinâmicas.

• O processo de geração é modelado como uma Equação Diferencial Estocástica (SDE) no espaço latente.
• Para evitar a necessidade de simulação cara da SDE durante o treinamento (o que seria computacionalmente proibitivo), os autores constroem uma aproximação da posterior variacional que permite a amostragem direta de $z_t$ sem simulação numérica.
• Utilizando a Transformada h de Doob e assumindo dinâmicas lineares com ruído aditivo, eles derivam uma densidade condicional $p(z_t|z_1, z_0)$ que é Gaussiana. Isso permite reparametrizar $z_t$ como uma combinação linear de $z_0$, $z_1$ e ruído $\epsilon$:
  $$z_t = \eta_t \epsilon + \kappa_t z_1 + \nu_t z_0$$
  Onde $\eta_t, \kappa_t, \nu_t$ são funções do tempo.

C. Função de Perda (Objetivo de Treinamento)

O objetivo de treinamento é maximizar o ELBO, que se traduz em minimizar uma perda composta por dois termos principais:

1. Termo de Reconstrução: $-\ln p_\theta(x_1|z_1)$, garantindo que o decodificador possa recuperar a imagem original.
2. Termo de Dinâmica (SI): Uma perda baseada na diferença entre a velocidade do processo variacional e a do modelo generativo no espaço latente.
   • A perda é formulada para ser livre de simulação (simulation-free), similar aos modelos de difusão no espaço de observação, mas aplicada às variáveis latentes.
   • Os autores introduzem uma parametrização chamada InterpFlow, que demonstrou ser mais estável e eficiente em termos de gradiente do que outras abordagens (como predição de ruído ou denoising direto).

D. Flexibilidade de Priori e Amostragem

• Priors Diversos: Diferente de modelos de difusão padrão que exigem priores Gaussianos, o LSI suporta priores arbitrárias ($p_0$), como Uniforme, Laplaciana ou Misturas Gaussianas, mantendo a flexibilidade do framework SI.
• Amostragem Flexível: O framework permite o uso de amostradores determinísticos (ODE de fluxo de probabilidade) ou estocásticos, além de suportar Class-Free Guidance (CFG) para controle condicional durante a geração.

3. Contribuições Principais

1. LSI (Latent Stochastic Interpolants): Um novo framework escalável para treinamento conjunto de modelos generativos de variáveis latentes com variáveis dinâmicas em tempo contínuo.
2. Perspectiva Unificadora: Integra a formulação flexível de tempo contínuo dos SI dentro de modelos de variáveis latentes, aproveitando insights de processos estocásticos contínuos.
3. Objetivo ELBO Principial: Uma nova função de perda ELBO que mantém as vantagens dos SI (treinamento livre de simulação, escolha flexível de prior) enquanto habilita o treinamento conjunto no espaço latente.
4. Eficiência Computacional: Demonstra que operar no espaço latente reduz drasticamente o custo computacional (FLOPs) durante a amostragem em comparação com modelos no espaço de pixels, sem sacrificar a qualidade da geração.

4. Resultados Experimentais

Os autores avaliaram o LSI no benchmark padrão de geração de imagens ImageNet em várias resoluções (64x64, 128x128, 256x256).

• Qualidade de Geração (FID): Os modelos LSI alcançaram FID (Fréchet Inception Distance) comparável aos modelos de SI treinados diretamente no espaço de observação (pixels), apesar de operarem em um espaço latente de dimensão reduzida.
  • Exemplo (128x128): LSI obteve FID de 3.12, comparável a modelos de observação com FID de 3.46.
• Eficiência (FLOPs e Parâmetros):
  • O modelo latente (L) possui menos parâmetros e requer menos FLOPs por passo de amostragem.
  • Como o codificador não é usado na amostragem e o decodificador é usado apenas uma vez, a economia computacional se acumula com o número de passos. Para 256x256, houve uma redução de 64.9% nos FLOPs de amostragem em comparação com o modelo de observação.
• Aprendizado Conjunto vs. Independente:
  • Experimentos mostraram que o treinamento conjunto (ajustando o peso $\beta$ entre reconstrução e geração) melhora significativamente o FID em comparação com codificadores pré-treinados ou treinados independentemente.
  • O modelo é robusto a mudanças na capacidade do modelo (transferência de camadas convolucionais do modelo latente para o codificador/decodificador), mantendo o desempenho mesmo com menos capacidade no modelo latente.
• Parametrização e Priors:
  • A parametrização InterpFlow superou outras (OrigFlow, NoisePred, Denoising).
  • O modelo funcionou bem com diferentes priores (Gaussiana, Uniforme, Laplaciana), confirmando a flexibilidade do framework.

5. Significado e Impacto

O trabalho Latent Stochastic Interpolants representa um avanço significativo na interseção entre modelos de difusão, VAEs e processos estocásticos contínuos.

• Superação de Limitações: Resolve a limitação histórica dos SI de não poderem ser aplicados em espaços latentes aprendidos, permitindo o uso de representações latentes eficientes com a flexibilidade de transformação de distribuições dos SI.
• Escalabilidade: Oferece uma alternativa viável e mais eficiente para a geração de imagens de alta resolução, reduzindo a carga computacional sem comprometer a fidelidade.
• Flexibilidade Teórica: Ao fornecer um controle de verossimilhança (via ELBO) e suportar priores arbitrárias, o LSI abre caminho para modelos generativos mais expressivos e adaptáveis, superando a dependência de priores simples (Gaussianas) comuns em VAEs e difusões tradicionais.

Em resumo, o LSI unifica a eficiência de modelos latentes com a flexibilidade teórica dos interpolantes estocásticos, estabelecendo um novo paradigma para a geração de dados complexos.