Resumo Técnico: Inversão de Transformações de Dados via Amostragem por Difusão

Declaração do Problema

O artigo aborda o problema da inversão de transformação em grupos de Lie gerais. Em muitos cenários de aprendizado de máquina e modelagem científica, os dados observados $\tilde{x}$ são gerados aplicando uma transformação desconhecida $g$ de um grupo de Lie $G$ a uma amostra dentro da distribuição $x$ (ou seja, $\tilde{x} = g \cdot x$). O objetivo é recuperar uma transformação inversa $g^{-1}$ que mapeie a observação fora da distribuição de volta para a distribuição de dados original.

Isso é formulado como um problema inverso cego porque a transformação de ruído é desconhecida e a solução é geralmente não única. Embora os métodos existentes frequentemente operem no espaço dos dados ou sejam especializados em famílias de transformações específicas (por exemplo, transformações afins ou grupos de rotação), este trabalho visa ações de grupos de Lie gerais, incluindo grupos não compactos, não abelianos e sem métricas bi-invariantes (ex: homografias de imagem $PGL(3, \mathbb{R})$ e simetrias de ponto de Lie para EDPs).

Uma aplicação fundamental destacada é a equivariância em tempo de teste. Redes neurais pré-treinadas frequentemente carecem de robustez a transformações geométricas. Os autores propõem inverter essas transformações em tempo de teste para trazer os inputs de volta para a distribuição de treinamento do modelo, melhorando assim a robustez sem a necessidade de arquiteturas equivariantes especializadas ou retreinamento.

Metodologia: Difusão de Energia para Inversão de Transformação (TIED)

O método proposto, TIED, modela a posterior sobre as transformações $p(g | \tilde{x})$ como uma distribuição de Boltzmann definida por uma função de energia $E$ no espaço dos dados. A energia é derivada de uma priori de dados $p_x(x)$, onde $E(x) = -\log p_x(x)$. A densidade posterior é dada por:
$$p(g | \tilde{x}) \propto p_x(g^{-1} \cdot \tilde{x}) |\det J_{g^{-1}}(\tilde{x})|$$

Para amostrar desta posterior, o TIED introduz um processo de difusão diretamente no grupo de Lie. A metodologia supera dois desafios principais: a natureza acidentada e multimodal das paisagens de energia (frequentemente derivadas de redes neurais) e as restrições geométricas de amostragem em variedades.

1. Difusão Trivializada em Grupos de Lie

Em vez de realizar atualizações nos espaços tangentes de elementos de grupo arbitrários, o TIED utiliza a trivialização. As atualizações são realizadas na álgebra de Lie $\mathfrak{g}$ (o espaço tangente na identidade) e transportadas para o elemento de grupo atual via o mapa exponencial.

• Processo de Avanço (Forward Process): Uma SDE (Equação Diferencial Estocástica) de explosão de variância é definida no grupo ao desenhar perturbações infinitesimais na álgebra de Lie e transportá-las via multiplicação à direita:
  $$dg_t = d(L_{g_t})_e [\gamma(t) dw^{\mathfrak{g}}_t]$$
  onde $w^{\mathfrak{g}}_t$ é o movimento browniano na álgebra de Lie.
• Processo de Recuo (Reverse Process): A SDE de tempo reverso é derivada para recuperar amostras da posterior. Crucialmente, o termo de deriva envolve o score trivializado $\nabla_g \log p_t(g_t)$, que é um elemento da álgebra de Lie.

2. Identidade de Score Alvo Trivializado

Uma contribuição teórica central é uma nova identidade de score alvo trivializado. O ajuste de score (score matching) padrão requer acesso a amostras limpas, que não estão disponíveis neste cenário de inversão. Em vez disso, o TIED expressa o score da distribuição ruidosa em termos do gradiente da função de energia limpa.
Os autores provam que para grupos de Lie gerais (sem exigir uma métrica bi-invariante), o score trivializado pode ser estimado via uma expectativa de Monte Carlo sobre o kernel de ruído $k_t$:
$$\nabla_g \log p_t(g_t) = \nabla_g \log \int_G k_t(w) \exp\left(-E(g_t w^{-1}) - \log \lambda(w)\right) d\mu(w)$$
onde $\lambda(w)$ contabiliza a mudança da medida de Haar sob inversão. Isso permite que o score seja estimado usando apenas a função de energia $E$ e seus gradientes, computados via diferenciação automática.

3. Implementação Prática

O algoritmo de amostragem (Algoritmo 1) discretiza a SDE de recuo. Em cada etapa, ele:

1. Amostra ruído $w$ do kernel de ruído $k_t$.
2. Computa a energia $E$ para candidatos transformados.
3. Estima o score trivializado usando a identidade de Monte Carlo (empregando o truque log-sum-exp para estabilidade).
4. Atualiza o elemento de grupo usando o mapa exponencial.

O método suporta várias definições de energia, incluindo aquelas derivadas de Autoencoders Variacionais (VAEs) ou diretamente da confiança de predição de um classificador pré-treinado.

Principais Contribuições

1. Amostrador TIED: Um novo amostrador de difusão para inverter transformações de dados em grupos de Lie gerais. Ele garante que as atualizações permaneçam na variedade e lida com ações não compactas e não lineares usando apenas o conhecimento da álgebra de Lie.
2. Estrutura Teórica: Uma prova de que a inversão de transformação pode ser usada para a equivariância em tempo de teste de modelos pré-treinados. Os autores demonstram que a posterior de inversão de transformação é $G$-equivariante, garantindo a solidez da abordagem.
3. Nova Identidade de Score: Uma generalização da identidade de score alvo para grupos de Lie que funciona sem métricas bi-invariantes, permitindo a amostragem eficiente baseada em score usando gradientes de energia em vez de amostras de dados limpos.
4. Validação Empírica: Demonstrações em homografias de imagem ($PGL(3, \mathbb{R})$) e simetrias de pontos de EDP ($SL(2, \mathbb{R}) \ltimes H(1, \mathbb{R})$), mostrando melhorias de desempenho consistentes sobre fortes baselines.

Resultados Experimentais

O artigo avalia o TIED em três tarefas principais:

1. Amostragem Sintética em $SO(10)$: O TIED amostra com sucesso de uma distribuição de energia multimodal em um grupo de Lie de alta dimensão ($dim=45$), convergindo significativamente mais rápido que um amostrador de Langevin cinético trivializado.
2. Classificação de Imagem (MNIST):
   • Configuração: Um ResNet18 pré-treinado é testado em imagens MNIST sujeitas a transformações afins ($Aff(2, \mathbb{R})$) e de homografia ($PGL(3, \mathbb{R})$) desconhecidas.
   • Baselines: Comparado contra métodos baseados em otimização (LieLAC, FoCal, ITS) e métodos de amostragem (Langevin Cinético).
   • Resultados: O TIED supera significativamente todos os baselines. Por exemplo, sob transformações afins com energia baseada em classificador, o TIED atinge 82,64% de acurácia comparado a 55,48% para o ResNet18 bruto e 74,72% para o Langevin Cinético. Ele também supera redes equivariantes especializadas (affConv) em algumas configurações.
3. Resolução de EDP:
   • Configuração: Um operador neural DeepONet é testado em equações de Calor e Burgers 1D sob transformações desconhecidas de simetria de ponto de Lie.
   • Resultados: O TIED atinge o menor erro relativo $L_2$ em todos os casos de teste, superando o LieLAC e o FoCal. Para a equação do Calor 1D com simetrias $SL(2, \mathbb{R}) \ltimes H(1, \mathbb{R})$, o TIED reduz o erro para 0,043, comparado a 0,066 para o LieLAC e 0,690 para o DeepONet não corrigido.

Significância e Alegações

Os autores afirmam que o TIED fornece um pipeline livre de treinamento para tornar qualquer modelo pré-treinado equivariante a transformações de grupos de Lie em tempo de teste. Ao alavancar modelos baseados em energia (que podem ser aproximados pelo próprio modelo pré-treinado), o método evita a necessidade de arquiteturas equivariantes especializadas ou retreinamento caro.

A significância reside em:

• Generalidade: Aplica-se a uma ampla classe de grupos de Lie, incluindo aqueles que são não compactos e carecem de métricas bi-invariantes, que têm sido desafiadores para métodos anteriores de difusão ou otimização.
• Robustez: Restaura efetivamente os inputs transformados para a distribuição de treinamento, melhorando significativamente a generalização de redes pré-treinadas para transformações geométricas fora da distribuição.
• Eficiência: A abordagem de difusão, combinada com a identidade de score trivializado, oferece um processo de amostragem mais eficiente e estável comparado ao MCMC tradicional ou otimização iterativa em paisagens de energia acidentadas.

O artigo conclui que esta abordagem representa uma forma de escalonamento de equivariância em tempo de inferência, oferecendo uma solução prática para robustez na era dos grandes modelos de fundação, onde lidar com transformações fora da distribuição é crítico.