Resumo Técnico: Reward Transport – Controle de Propriedades em Flow Matching via Alinhamento no Espaço de Ruído

1. Declaração do Problema

No Flow Matching (FM) e em modelos generativos relacionados baseados em difusão, o acoplamento (coupling) — a regra que determina qual vetor de ruído é pareado com qual ponto de dado durante o treinamento — é tipicamente tratado como um heurístico computacional. Abordagens padrão usam pareamento independente por simplicidade ou transporte ótimo (OT) de minibatch para retificar trajetórias e reduzir a variância. Nesses paradigmas, o acoplamento é otimizado puramente para facilitar o treinamento, não para codificar propriedades estruturais específicas no campo de fluxo aprendido.

Consequentemente, alcançar a geração controlável (por exemplo, gerar moléculas com scores específicos de logP ou QED) geralmente requer mecanismos explícitos no momento da inferência, tais como:

• Classificador sem guia (Classifier-free guidance): Requer um valor de propriedade alvo como entrada.
• Guia por gradiente (Gradient guidance): Computa gradientes de um preditor de propriedade em cada etapa de amostragem.
• Aprendizado por Reforço (RL): Ajusta o gerador contra um oráculo, muitas vezes exigindo consultas ao oráculo durante a amostragem.

Esses métodos introduzem overhead de inferência, dependência de modelos de recompensa ou oráculos externos, ou pipelines de treinamento complexos. O artigo questiona se o próprio acoplamento pode servir como uma interface de alinhamento, incorporando estrutura controlável diretamente no campo de fluxo durante o treinamento, de modo que o controle esteja disponível na inferência sem computação adicional ou chamadas de oráculo.

2. Metodologia: Reward Transport

Os autores propõem o Reward Transport, um método que redefine o acoplamento ruído-dado como um mecanismo para injetar estrutura de tarefa no campo de fluxo.

2.1 Acoplamento de Transporte Ótimo Alinhado por Propriedade

Em vez de parear ruído e dados de forma independente, o Reward Transport constrói um acoplamento monotônico baseado em uma propriedade molecular alvo $y(x)$ (ex: logP, QED).

1. Chaves Escalares: Para cada vetor de ruído $z$ (uma sequência de embeddings de tokens), uma chave de ordenação escalar $s(z)$ é computada. A escolha primária é a norma euclidiana da média agrupada (mean-pooled): $s(z) = \|\bar{z}\|_2$.
2. Pareamento Baseado em Ranking: Dado um lote de vetores de ruído $\{z_n\}$ e moléculas $\{x_n\}$ com propriedades $\{y_n\}$, ambos os conjuntos são ordenados por suas respectivas chaves ($s$ e $y$). Eles são então pareados rank-a-rank: o vetor de ruído com o $k$-ésimo menor $s$ é pareado com a molécula que possui o $k$-ésimo menor $y$.
3. Objetivo de Treinamento: O modelo de flow matching $v_\theta$ é treinado para transportar esses pares alinhados. A perda combina um MSE padrão sobre a predição do endpoint (em espaço de embedding contínuo) e uma perda de Entropia Cruzada (CE) para garantir que os tokens decodificados formem moléculas válidas.

2.2 Condicionamento de Direção

Para garantir que o modelo possa acessar a coordenada do espaço de ruído $s$ durante a inferência, um MLP de Direção é introduzido. Este MLP de duas camadas mapeia o escalar normalizado $s$ para um embedding de dimensão $d$, que é adicionado à entrada do transformer junto com os embeddings de tempo e posição. Isso permite que o modelo condicione seu campo de velocidade à coordenada de ruído específica usada no pareamento monotônico.

2.3 Mecanismo de Inferência

Na inferência, o usuário especifica um sinal de direção escalar $s^*$ (uma coordenada do espaço de ruído, não um valor de propriedade alvo). O modelo gera moléculas integrando o campo de fluxo partindo de vetores de ruído onde $s(z) \geq s^*$.

• Garantia Teórica (Proposição 1): No limite onde o fluxo aprendido preserva perfeitamente a correspondência de ranking entre $s$ e $y$ (limite de preservação de acoplamento), o condicionamento em $s \geq s^*$ induz uma distribuição equivalente a um passo de seleção do Método de Entropia Cruzada (CEM). Especificamente, ele trunca a distribuição gerada para a fração top-$\epsilon$ da distribuição da propriedade, onde $\epsilon$ é determinado pelo limiar $s^*$.
• Zero Overhead: Este controle não requer oráculo, não requer modelo de recompensa, não requer computação de gradiente e não requer amostragem de rejeição no momento da inferência.

3. Principais Contribuições

1. Acoplamento como Interface de Alinhamento: O artigo reformula o acoplamento de ruído-dado no flow matching de um heurístico de treinamento para uma interface de alinhamento de nível de distribuição. Demonstra que escolher o acoplamento baseado em uma propriedade alvo organiza o campo de fluxo aprendido de tal forma que as coordenadas do espaço de ruído correspondem aos valores das propriedades.
2. Direcionamento Molecular de Botão Único (Single-Knob): O método alcança controle monotônico sobre propriedades moleculares (logP e QED) usando um único escalar de ruído. Crucialmente, o mesmo botão produz respostas estruturais opostas para diferentes alvos: aumentar $s$ gera moléculas maiores para logP (lipofilicidade), mas moléculas menores para QED (drug-likeness). Isso descarta um viés genérico de tamanho e confirma que o acoplamento escreve a estrutura específica do alvo no fluxo.
3. Reward Transport para Sequências Discretas: Os autores implementam isso para sequências de tokens SELFIES de comprimento variável, identificando e abordando dois modos de falha específicos em modelos de fluxo que geram dados discretos:
   • Apagamento Implícito de Sinal: Em arquiteturas Pre-LayerNorm, a informação de magnitude (codificada em $\|\bar{z}\|$) é apagada. O MLP de Direção resolve isso injetando um padrão de direção aprendido.
   • Confundimento de Comprimento (Length Confound): Sem a perda MSE não mascarada nos tokens de padding, o modelo aprende um atalho onde o alto $s$ simplesmente aumenta o comprimento da sequência para satisfazer a restrição de norma. O MSE não mascarado força o modelo a regredir o ruído bruto nas posições de padding, ancorando o sinal de magnitude e prevenindo esse atalho.

4. Resultados Experimentais

O método foi avaliado nos datasets ZINC-250K e GuacaMol.

• Desempenho de Controle:
  • logP: No ZINC-250K, percorrer $s$ de -3 a +7 deslocou o logP médio de 1.50 para 5.44 (+137%), com uma correlação de Spearman $\rho=1.000$ entre $s$ e a propriedade média. O tamanho do efeito (Cohen's $d$) atingiu 1.63.
  • QED: O mesmo botão deslocou o QED monotonicamente (ex: +0.059 em $s=5$ no ZINC), com $\rho=1.000$.
  • Especificidade Estrutural: Para logP, aumentar $s$ aumentou o número de átomos pesados (12 $\to$ 23). Para QED, aumentar $s$ diminuiu o número de átomos pesados (23 $\to$ 16).
• Generalização: Os resultados foram replicados no GuacaMol (1.6M moléculas), mostrando efeitos ainda mais fortes para logP ($d=2.31$) e monotonicidade consistente para QED dentro do intervalo válido.
• Comparação com Baselines:
  • Comparado ao Conditional Flow Matching (onde os valores de propriedade são entradas explícitas), o Reward Transport alcança monotonicidade semelhante ($\rho=1.000$), mas opera via truncamento de distribuição em vez de atingir um alvo explícito.
  • Comparado à difusão de predição de $\epsilon$ ($\epsilon$-prediction diffusion), o método mostra que o alinhamento ao nível de acoplamento está estruturalmente ausente na predição de $\epsilon$, pois o alvo Bayes-ótimo é independente do acoplamento, atenuando o sinal de gradiente.
• Dinâmica de Treinamento: A preservação de ranking (a correlação entre $s$ e o $y$ gerado) emerge rapidamente durante o ajuste fino (dentro de um epoch) e estabiliza abaixo de 1.0 (ex: 0.57 para logP), consistente com a lacuna de "aproximação de distribuição" descrita na Proposição 1.

5. Significância e Alegações

O artigo afirma que o Reward Transport estabelece o acoplamento no flow matching como um canal de controle distinto, de nível de distribuição, que pode carregar a estrutura de alinhamento por si só.

• Eficiência: Elimina a necessidade de chamadas de oráculo, modelos de recompensa ou otimização iterativa no momento da inferência, reduzindo significativamente o custo computacional em comparação com métodos de RL ou de guia.
• Mecanismo: Demonstra que a "direcionabilidade" de modelos generativos não precisa depender de entradas de condicionamento explícitas; ela pode ser integrada à geometria do modelo via o acoplamento de treinamento.
• Limitações e Escopo: Os autores são modestos quanto ao escopo. O método fornece um botão distributivo de 1D (truncando a distribuição) em vez de uma atribuição de valor exato por amostra. Funciona bem para propriedades escalares correlacionadas com características moleculares globais (como tamanho ou polaridade), mas falha para propriedades topológicas como a Acessibilidade Sintética (SA), onde uma chave escalar 1D não consegue distinguir estruturas complexas de tamanho similar. Além disso, o método mostra ser incompatível com alvos de predição de $\epsilon$, destacando que a interface depende do alvo de predição da rede.

Em resumo, o trabalho desloca o paradigma da geração controlável de "guiar a trajetória na inferência" para "alinhar a distribuição de treinamento", ofereando um mecanismo de controle de propriedade de zero-overhead e baseado em princípios através do acoplamento.