Variance Reduction for Non-Log-Concave Sampling with Applications to Inverse Problems

Resumo Técnico: Redução de Variância para Amostragem Não-Log-Côncava com Aplicações a Problemas Inversos

Enunciado do Problema
A amostragem de distribuições de alta dimensão, não-log-côncavas, com densidades não normalizadas, é um desafio fundamental em aprendizado de máquina, particularmente quando os gradientes exatos da função potencial não estão disponíveis. Nesses cenários, os gradientes devem ser aproximados via estimativas estocásticas, que frequentemente exibem alta variância sob um orçamento computacional fixo por iteração. Embora técnicas de redução de variância (ex: SGD com momentum, STORM, PAGE) tenham melhorado a convergência em otimização não-convexa, sua aplicação à amostragem de distribuições não-log-côncavas permanece amplamente inexplorada. Além disso, os métodos existentes de amostragem posterior para problemas inversos utilizando modelos generativos baseados em score (SGM) operam tipicamente em um regime de lote completo (full-batch), limitando sua aplicabilidade a conjuntos de dados de grande escala comuns em imagens 3D e dinâmicas.

Metodologia
Os autores propõem um framework unificado que adapta técnicas de redução de variância da otimização não-convexa para o domínio da amostragem. Eles introduzem dois estimadores unificados:

1. Difusão de Langevin com Redução de Variância de Lipschitz Médio (ML-VRLD): Este estimador unifica o SGD com momentum e seu correspondente estocástico. Ele assume a continuidade de Lipschitz dos gradientes exatos (Assunção 1) e fornece um tratamento unificado onde o estimador é uma média ponderada de um gradiente de mini-lote (mini-batch) e uma atualização baseada em momentum.
2. Difusão de Langevin com Redução de Variância de Lipschitz de Amostra (SL-VRLD): Este estimador unifica PAGE e STORM. Ele assume a continuidade de Lipschitz dos gradientes estocásticos no sentido de média quadrática (Assunção 4). Ele incorpora um termo de correção baseado na diferença entre os gradientes estocásticos atuais e passados, o que melhora a precisão da estimativa sob a condição de Lipschitz de amostra.

Ambos os métodos são projetados para alcançar a convergência de Informação de Fisher (FI) relativa $\epsilon$ com $O(1)$ computações de gradiente estocástico por iteração, contrastando com o requisito de $O(Ld\sigma^2/\epsilon)$ do SGLD padrão.

O framework é estendido para problemas inversos (ex: reconstrução de MRI e CT) integrando estes estimadores em um processo de amostragem de Monte Carlo de Langevin (LMC) com priors SGM. A formulação de amostragem posterior utiliza uma sequência de distribuições alvo intermediárias com cronogramas de recozimento (annealing) e suavização para guiar as iterações de um prior suavizado para o posterior verdadeiro. Garantias teóricas são estabelecidas sob suposições de erro de estimativa SGM limitadas.

Principais Contribuições

• Análise Unificada: O artigo fornece a primeira análise unificada de estimadores de redução de variância (especificamente ML-VRLD e SL-VRLD) para amostragem de distribuições não-log-côncavas.
• Complexidade Melhorada:
  • Tanto ML-VRLD quanto SL-VRLD alcançam convergência de FI relativa $\epsilon$ com $O(1)$ computações de gradiente por iteração, melhorando significativamente a complexidade por iteração do SGLD.
  • O SL-VRLD reduz ainda mais a complexidade computacional total de $O(L^3d^3\sigma^2/\epsilon^3)$ (SGLD) para $O(L^2d^2\sigma^2/\epsilon^3)$.
• Garantias de Convergência: Os autores estabelecem taxas de convergência não-assintóticas em distância de variação total (TV) quadrada sob uma suposição de desigualdade de Poincaré e provam a convergência fraca para a distribuição alvo para ambos os algoritmos.
• Extensão para Problemas Inversos: O framework é estendido para amostragem posterior com priors SGM, estabelecendo as mesmas garantias de convergência sob erro de estimativa SGM limitado. Isso aborda a limitação dos métodos de amostragem posterior de lote completo existentes.

Resultados

• Experimentos Sintéticos: Em modelos de mistura Gaussiana (GMM) e problemas de compressão de sinais, os métodos propostos demonstram convergência mais rápida para valores de FI relativa mais baixos em comparação ao SGLD. O SGLD com um tamanho de lote fixo converge para um nível de FI suboptimal, enquanto o ML-VRLD e o SL-VRLD alcançam a precisão alvo com o mesmo tamanho de lote fixo. O SL-VRLD consistentemente atinge a distribuição alvo mais rápido que o ML-VRLD.
• Validação Estatística: Em uma tarefa de amostragem posterior com um ground truth conhecido, o SGLD produziu amostras com variância significativamente inflada. O ML-VRLD e o SL-VRLD reduziram significativamente essa variância, com o SL-VRLD alcançando a correspondência mais próxima da variância do posterior do ground truth.
• Aplicações de Imagem:
  • Reconstrução de MRI: No conjunto de dados fastMRI com fatores de aceleração de 8 e 12, o ML-VRLD e o SL-VRLD produziram detalhes anatômicos mais nítidos e menos artefatos comparados aos baselines (incluindo DPS, PnP-DM e SGLD). O SL-VRLD alcançou as maiores pontuações de PSNR e SSIM.
  • Reconstrução de CT de Ângulo Esparso: No conjunto de dados 2DeteCT com 180 e 360 vistas de projeção, os métodos propostos superaram os baselines na preservação de detalhes estruturais finos, demonstrando o benefício da redução de variância em regimes de alto ruído e subamostragem.

Significância
O artigo afirma fornecer as primeiras garantias de convergência não-assintótica para amostragem não-log-côncava usando estimadores de redução de variância com $O(1)$ computações de gradiente por iteração. Ao estender essas técnicas para problemas inversos com priors SGM, este trabalho possibilita a amostragem posterior fundamentada para tarefas de imagem de larga escala onde métodos de lote completo são computacionalmente proibitivos. Os resultados empíricos validam que a redução de variância consistentemente melhora a qualidade da amostra e a precisão estatística sob restrições fixas de memória e computação, oferecendo uma solução escalável para inferência Bayesiana de alta dimensão.