🤖 machine learning

Riemannian MeanFlow

Este artigo apresenta o Riemannian MeanFlow (RMF), um framework que aprende mapas de fluxo direto em variedades para permitir modelagem generativa de alta qualidade com até 10 $\times$ menos avaliações de função do que os modelos de difusão ou fluxo tradicionais, ao mesmo tempo que facilita o projeto guiado por recompensa eficiente.

Autores originais: Dongyeop Woo, Marta Skreta, Seonghyun Park, Kirill Neklyudov, Sungsoo Ahn

Publicado 2026-05-04

📖 5 min de leitura🧠 Leitura aprofundada

CC BY 4.0

Autores originais: Dongyeop Woo, Marta Skreta, Seonghyun Park, Kirill Neklyudov, Sungsoo Ahn

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você está tentando ensinar um robô a projetar novas proteínas ou sequências de DNA. Esses não são apenas listas de letras ou números; são estruturas 3D complexas com regras estritas, como um quebra-cabeça onde cada peça deve encaixar perfeitamente em uma forma geométrica específica.

No mundo da IA, os modelos de "Difusão" e "Fluxo" são como artistas que transformam lentamente um borrão de ruído em uma obra-prima perfeita. Geralmente, eles fazem isso dando centenas de passos minúsculos e cuidadosos, verificando seu trabalho a cada passo único para garantir que não se desviem do caminho.

O Problema: A Caminhada Lenta
O artigo aponta um gargalo majoritário: dar centenas de passos é incrivelmente lento e caro. É como tentar atravessar uma sala dando 100 passos minúsculos e cautelosos quando você poderia simplesmente atravessar em uma ou duas passadas largas. Para cientistas que precisam projetar milhares de moléculas, essa lentidão é um impeditivo.

A Solução: Riemannian MeanFlow (RMF)
Os autores introduzem um novo método chamado Riemannian MeanFlow (RMF). Pense nisso como ensinar o robô a "ver o caminho inteiro" de uma só vez, em vez de apenas o próximo passo.

Veja como eles fizeram isso, usando algumas analogias do cotidiano:

1. O Atalho da "Velocidade Média"

Imagine que você está dirigindo da Cidade A para a Cidade B.

Jeito Antigo (Difusão/Fluxo): Você dirige devagar, verificando seu GPS a cada segundo, ajustando o volante constantemente. Isso é preciso, mas leva uma eternidade.
O Jeito Novo (RMF): Em vez de verificar a cada segundo, você calcula a velocidade média e a direção necessárias para ir de A a B de uma só vez. Você não precisa parar e verificar; apenas dirige o caminho inteiro baseado nesse único cálculo.

O artigo prova matematicamente que é possível aprender essa "velocidade média" diretamente nas formas geométricas complexas (variedades) onde essas moléculas vivem, sem necessidade de dar centenas de passos minúsculos.

2. Três Maneiras de Descrever a Jornada

Os autores descobriram três maneiras diferentes de descrever essa "velocidade média" que são todas matematicamente equivalentes. Elas são como descrever uma viagem de estrada a partir de três perspectivas diferentes:

A Visão Euleriana: Olhar para a estrada a partir de um ponto fixo (como uma câmera em um poste) e ver como o tráfego flui passando por você.
A Visão Lagrangiana: Sentar dentro do carro e sentir como a velocidade muda enquanto você dirige.
A Visão Semigrupo: Dividir a viagem em duas partes (de A a B, depois de B a C) e perceber que o caminho direto de A a C é apenas a combinação dessas duas partes.

O artigo mostra que usar essas perspectivas permite que a IA aprenda o salto de "um passo" diretamente.

3. O Truque do "Stop-Gradient"

Treinar esses modelos é como tentar equilibrar uma pilha de pratos enquanto faz malabarismo. Se você tentar calcular como cada mudança minúscula afeta o resultado final, a matemática fica tão pesada que o computador trava (ou fica sem memória).
Os autores usam um truque inteligente chamado "stop-gradient". Imagine que você está ensinando um aluno. Em vez de fazê-lo recalcular toda a história de como chegou a uma resposta cada vez que ele erra, você apenas diz: "Aqui está a resposta correta, agora apenas faça sua próxima tentativa combinar com ela". Isso impede que o computador fique sobrecarregado com matemática complexa, tornando o treinamento mais rápido e estável.

4. Prever o Destino, Não a Velocidade

A maioria dos modelos tenta prever a velocidade a cada momento. O RMF tenta prever o destino (o ponto final no mapa) diretamente.

Analogia: Em vez de dizer a um motorista: "Vire à esquerda a 8 km/h, depois vire à direita a 16 km/h", você apenas diz: "Dirija até ver o Celeiro Vermelho Grande".
Essa abordagem funciona muito melhor para dados de alta dimensão (como longas sequências de DNA ou proteínas grandes) e permite que o modelo use software existente projetado para essas formas sem necessidade de reescrita.

O Que Eles Conquistaram
O artigo testou isso em duas tarefas científicas do mundo real:

Projetar Promotores de DNA: Estes são como "botões de início" para genes. O novo método criou sequências de DNA de alta qualidade em um passo que foram tão boas quanto os métodos antigos que levavam 100 passos.
Projetar Esqueletos de Proteínas: Estes são os esqueletos das proteínas. O novo método pôde gerar estruturas de proteínas válidas em 1 a 5 passos, enquanto métodos anteriores colapsavam ou falhavam completamente quando solicitados a fazer isso tão rápido.

A Conclusão
O artigo afirma que, ao entender a geometria dos dados e usar esses atalhos de "velocidade média", eles podem gerar estruturas biológicas complexas 10 vezes mais rápido (requerendo 10 vezes menos cálculos de computador) sem perder qualidade. Isso torna possível usar IA para design científico rápido e iterativo, onde pesquisadores podem testar rapidamente milhares de ideias em vez de esperar horas por um único resultado.

Resumo Técnico: Riemannian MeanFlow

Declaração do Problema

Os modelos de difusão e fluxo tornaram-se o paradigma dominante para modelagem generativa em variedades Riemannianas, sendo aplicados com sucesso a tarefas como geração de espinha dorsal de proteínas e projeto de sequências de DNA. No entanto, esses métodos geralmente exigem dezenas a centenas de avaliações de rede neural (avaliações de função, ou NFEs) durante a inferência para integrar numericamente equações diferenciais ordinárias (EDOs) ou equações diferenciais estocásticas (EDEs). Esse ônus computacional cria um gargalo, particularmente em pipelines de projeto científico onde os modelos generativos servem como mecanismos de proposta dentro de loops iterativos para otimização guiada por propriedades. Embora técnicas de geração em poucos passos, como modelos de consistência e aprendizado de mapas de fluxo, tenham sido exploradas em contextos euclidianos, elas permanecem pouco exploradas e frequentemente instáveis em configurações de variedades Riemannianas de alta dimensão.

Metodologia: Riemannian MeanFlow (RMF)

Os autores introduzem o Riemannian MeanFlow (RMF), um framework projetado para aprender mapas de fluxo diretamente em variedades, permitindo geração de alta qualidade com tão poucos quanto um único passe forward. O núcleo da metodologia envolve derivar e utilizar três caracterizações equivalentes da velocidade média da variedade, que relacionam o mapa de fluxo entre pontos temporais arbitrários $s$ e $t$ .

1. Fundamentos Teóricos

O framework define a velocidade média $u_{s,t}$ como a velocidade constante que, se mantida ao longo de uma geodésica de $x_s$ a $x_t$ durante o tempo $t-s$ , transportaria o ponto corretamente. Os autores derivam três identidades que servem como base para objetivos de treinamento:

Identidade Euleriana: Derivada diferenciando a relação definidora em relação ao tempo fonte $s$ . Relaciona a velocidade média à derivada covariante do campo de velocidade ao longo da curva integral e à derivada do mapa logarítmico.
Identidade Lagrangiana: Derivada diferenciando em relação ao tempo alvo $t$ . Relaciona a velocidade média ao diferencial do mapa logarítmico aplicado à velocidade instantânea e à derivada parcial temporal da velocidade média.
Identidade de Semigrupo: Uma condição algébrica derivada da propriedade de semigrupo dos mapas de fluxo ( $\Phi_{r,t} \circ \Phi_{s,r} = \Phi_{s,t}$ ). Impõe consistência entre o fluxo direto de $s$ para $t$ e a composição de fluxos através de um tempo intermediário $r$ .

2. Objetivos de Treinamento e Estabilização

Os autores convertem essas identidades em objetivos de regressão. Uma escolha de design crítica é a colocação do operador stop-gradient. Diferentemente de trabalhos anteriores (por exemplo, Mapas de Fluxo Generalizados) que aplicam stop-gradients parcialmente, o RMF aplica o stop-gradient a todo o alvo de regressão. Essa formulação:

Elimina a necessidade de avaliar derivadas de ordem superior (por exemplo, gradientes através de produtos vetor-Jacobiano) durante a retropropagação.
Reduz o uso de memória em quase 2 $\times$ .
Melhora a estabilidade do treinamento em configurações de alta dimensão.

Para abordar dificuldades de otimização em altas dimensões, o artigo propõe várias técnicas de estabilização:

Amostragem Temporal: Pares de tempo ordenados ( $s \le t$ ) para objetivos Eulerianos e de Semigrupo, e pares não ordenados para objetivos Lagrangianos.
Ponderação Adaptativa da Perda: Reponderação da perda com base na magnitude do alvo de regressão para lidar com alta variância em termos dependentes de derivadas.
Controle de Derivada Temporal: Uso de incorporações temporais de baixa frequência para limitar a magnitude das derivadas temporais na saída da rede.

3. Parametrização

O artigo avalia três esquemas de parametrização para a velocidade média:

Previsão de $v$ : Prever o vetor tangente (velocidade).
Previsão de $x_t$ : Prever o ponto final diretamente.
Previsão de $x_1$ : Prever o ponto final da trajetória em $t=1$ dado o estado atual e o intervalo de tempo.
Os autores encontram que a previsão de $x_1$ é a mais robusta para variedades de alta dimensão. É compatível com arquiteturas científicas existentes (por exemplo, aquelas que emitem quadros SE(3) para proteínas) e escala bem para dimensões até $D=2048$ , enquanto a previsão de $x_t$ falha e a previsão de $v$ degrada-se à medida que a dimensionalidade aumenta.

4. Inferência Guiada por Recompensa

O framework suporta geração guiada por recompensa sem re-treinamento do modelo. Ao aproveitar a capacidade do mapa de fluxo de prever estados terminais a partir de passos intermediários, os autores propõem uma estratégia de olhar adiante de recompensa (reward look-ahead). Em vez de usar o gradiente da recompensa no estado ruidoso atual $x_t$ , eles usam o gradiente no ponto final previsto $\Phi_{t,1}(x_t)$ . Essa abordagem mostrou-se mais eficaz para direcionar gerações em direção a objetivos específicos.

Resultados Chave

Os autores validam o RMF em duas tarefas biológicas: Projeto de DNA Promotor (sequências em um simplex em $\mathbb{R}^{1024 \times 4}$ ) e Geração de Espinha Dorsal de Proteínas (em $SE(3)^N$ com $N$ até 128).

Eficiência: O RMF alcança qualidade de amostra comparável a baselines multi-passos de última geração (por exemplo, Fisher FM, FrameFlow, FrameDiff), exigindo até 10 $\times$ menos avaliações de função.
- Na geração de espinha dorsal de proteínas, o RMF alcança 82% de projetabilidade em 5 passos e 35% em 1 passo, enquanto as baselines caem abruptamente ou colapsam completamente nesses regimes de poucos passos.
- No projeto de DNA, variantes do RMF igualam o desempenho de 100 passos do Fisher FM com apenas 1 passo.
Estabilidade: As técnicas de estabilização propostas (especificamente a combinação de objetivos de Semigrupo e previsão de $x_1$ ) permitem treinamento estável em espaços de alta dimensão onde outros métodos de mapa de fluxo (como GFM) falham em convergir ou produzem amostras de alto erro.
Orientação por Recompensa: A inferência guiada por recompensa usando a estratégia de olhar adiante de $x_1$ melhora significativamente o alinhamento com perfis regulatórios alvo (DNA) e composições de estrutura secundária (proteínas) em comparação com orientação ingênua ou geração não orientada.

Significado e Alegações

O artigo afirma estabelecer uma fundamentação teórica principial e diretrizes práticas para amostragem rápida em variedades Riemannianas. Seu significado reside em:

Preenchendo a Lacuna: Estende capacidades de geração em poucos passos de espaços euclidianos para variedades Riemannianas complexas e de alta dimensão, um passo crítico para aplicações científicas.
Escalabilidade: Ao derivar identidades que evitam derivadas de ordem superior e utilizar previsão de $x_1$ , o método escala para dimensões ( $D \approx 2000$ ) e comprimentos de sequência anteriormente inacessíveis a modelos generativos de variedades em poucos passos.
Eficiência em Fluxos de Trabalho Científicos: A redução no custo de inferência (até 10 $\times$ ) torna os modelos generativos mais viáveis como mecanismos de proposta em loops iterativos de projeto científico, como otimização guiada por propriedades para descoberta de fármacos ou biologia sintética.
Olhar Adiante de Recompensa: Demonstra que mapas de fluxo permitem projeto guiado por recompensa eficiente ao prever estados terminais a partir de passos intermediários, uma capacidade não mostrada anteriormente em variedades.

Os autores posicionam o RMF não como um substituto para toda a modelagem generativa, mas como um framework especializado que permite geração eficiente e de alta qualidade em domínios científicos geometricamente restritos onde o custo computacional é um fator limitante.