Resumo Técnico: Velocity Scheduled Flow Matching (VSFM)

Declaração do Problema

O flow matching é um arcabouço de modelagem generativa em tempo contínuo onde uma rede neural $v_\theta(x, t)$ é treinada para regredir a velocidade condicional ao longo de um interpolante linear entre o ruído $x_0$ e os dados $x_1$. O flow matching padrão fixa o interpolante como a linha reta $x_t = x_0 + t(x_1 - x_0)$, o que implicitamente compromete o amostrador a mover-se a uma velocidade constante ao longo da trajetória.

O custo de amostragem é determinado pelo Número de Avaliações de Função (NFE) necessárias para integrar a Equação Diferencial Ordinária (ODE) $dx/dt = v_\theta(x, t)$. No regime de baixo NFE (ex: 4–50 passos), a trajetória deve ser percorrida eficientemente usando integradores simples como o Euler. O artigo postula que a suposição de velocidade constante é uma escolha de design arbitrária que pode não estar alinhada com as características de erro de discretização do campo de velocidade aprendido, particularmente próximo à variedade de dados (data manifold), onde o campo varia mais rapidamente.

Metodologia

1. Alvo de Treinamento Generalizado

Os autores introduzem o Velocity Scheduled Flow Matching (VSFM), que generaliza o alvo condicional de um deslocamento constante para uma forma variante no tempo:
$$u_t(x | x_0, x_1) = v(t) \cdot (x_1 - x_0)$$
onde $v: [0, 1] \to \mathbb{R}_{\ge 0}$ é um perfil de velocidade não negativo que satisfaz $\int_0^1 v(t) dt = 1$. Esta restrição garante que o deslocamento total permaneça $x_1 - x_0$, enquanto $v(t)$ controla a taxa de travessia.

O objetivo de treinamento regride a rede sobre o deslocamento de magnitude constante $(x_1 - x_0)$ em posições $x_t$ determinadas pelo perfil de movimento $\alpha(t) = \int_0^t v(s) ds$:
$$L_{VSFM}(\theta) = \mathbb{E}_{t, x_0, x_1} \left[ \| v_\theta(x_t, t) - (x_1 - x_0) \|^2 \right]$$
Crucialmente, o perfil de velocidade $v(t)$ influencia o treinamento apenas através da distribuição das posições $x_t$, e não pela ponderação da magnitude da perda. Isso evita instabilidade numérica associada ao peso de $v(t)^2$ próximo a $t=1$.

2. Perfis de Velocidade Polinomiais

O estudo instancia o VSFM com seis perfis polinomiais derivados de problemas de valor de contorno de planejamento de movimento:

• Baseline: Linear ($v(t)=1$).
• Subfamília de Frenagem (Braking): Velocidade monotonicamente decrescente (Desaceleração uniforme, Coasting, Pouso suave). Estes movem-se rápido perto de $t=0$ e desaceleram até o repouso em $t=1$.
• Subfamília Simétrica: Velocidade zero em ambos os extremos (Ease-in-out, Minimum-jerk).

3. Reescalonamento no Tempo de Inferência (Sem Retreinamento)

Uma contribuição fundamental é a capacidade de aplicar qualquer perfil admissível a um modelo de flow matching linear pré-treinado sem necessidade de retreinamento.

• Mecanismo: O modelo pré-treinado $v_\theta$ é treinado na trajetória linear onde $x_\tau = x_0 + \tau(x_1 - x_0)$. Para amostrar sob um novo perfil $v(t)$, os autores definem um campo adaptado $\tilde{v}(x, t) = v(t) v_\theta(x, \alpha(t))$.
• Implementação: Ao mudar variáveis para $\tau = \alpha(t)$, a ODE torna-se $dx/d\tau = v_\theta(x, \tau)$. O amostrador integra esta ODE em uma grade não uniforme $\tau_i = \alpha(i/N)$.
• Algoritmo: O passo de atualização é $x_{i+1} = x_i + v_\theta(x_i, \tau_i) \cdot \Delta \tau_i$, onde $\Delta \tau_i = \alpha((i+1)/N) - \alpha(i/N)$. Isso efetivamente executa o amostrador Euler padrão em uma grade de tempo não uniforme induzida pelo perfil.

4. Análise de Erro de Discretização

O artigo atribui os ganhos de desempenho ao erro de truncamento local do integrador Euler. O erro no passo $i$ escala com $(\Delta \tau_i)^2 \cdot \frac{d^2x}{d\tau^2}$.

• O termo $\frac{d^2x}{d\tau^2}$ (variação do campo) é maior próximo a $\tau=1$ (a variedade de dados) onde o campo de velocidade é altamente não linear.
• Perfis de frenagem atribuem pequenos $\Delta \tau_i$ (passos pequenos) perto de $t=1$, thereby suprimindo a principal fonte de erro.
• Perfis simétricos concentram grandes passos no ponto médio, o que é menos eficaz para a geração de ruído-para-dados, onde a não linearidade crítica está no ponto final.

Principais Resultados

Experimentos foram conduzidos no CIFAR-10 (32x32) usando um UNet condicional de classe e um integrador Euler de passo fixo.

1. Treinamento do Zero

O treinamento sob o perfil Coasting apresentou os melhores resultados:

• 4 NFE: FID reduzido em 17,4% (35,39 vs. 42,86 baseline).
• 50 NFE: FID reduzido em 10,4% (8,29 vs. 9,25 baseline).
• A desaceleração uniforme também melhorou em relação ao baseline, enquanto o Pouso suave (Smooth landing) teve desempenho inferior em baixo NFE devido aos passos iniciais excessivamente grandes.
• Perfis simétricos mostraram ganhos marginais em alto NFE, mas ficaram atrás do baseline em baixo NFE.

2. Adaptação no Tempo de Inferência (Sem Retreinamento)

Aplicando o perfil Coasting ao checkpoint público do TorchCFM (treinado linearmente) sem retreinamento:

• 50 NFE: FID reduzido em 19,8% (4,38 vs. 5,47 baseline).
• 32 NFE: FID reduzido em 18,6% (5,23 vs. 6,43 baseline).
• 100 NFE: A desaceleração uniforme alcançou o melhor FID (3,89, -16,0%).
• O Pouso suave inicialmente performou pior em 32 NFE (7,43 vs. 6,43) devido aos grandes passos iniciais introduzirem erro de truncamento, mas melhorou em orçamentos de NFE mais altos.

Significância e Alegações

O artigo afirma que o perfil de velocidade é uma alavanca simples e de baixo custo para melhorar a amostragem de poucos passos em flow matching. A principal significância reside em duas áreas:

1. Adaptação Sem Retreinamento: Demonstra que um modelo de flow matching linear pré-treinado pode ser otimizado para orçamentos específicos de NFE simplesmente mudando o esquema de integração (o perfil de velocidade), sem exigir treinamento adicional ou mudanças arquiteturais.
2. Insight Teórico: Identifica que a suposição de velocidade constante no flow matching padrão é subótima para a integração Euler. Ao alinhar o tamanho do passo com o erro de truncamento local (tomando passos menores onde o campo varia mais), ganhos significativos de qualidade são alcançados, particularmente no regime de baixo NFE.

Os autores observam que estas descobertas são específicas para o integrador Euler e para a tarefa de geração ruído-para-dados. Eles sugerem que os benefícios podem diferir para tarefas de imagem-para-imagem (onde a origem é estruturada) ou quando se utilizam solvers de ordem superior que cancelam erros de truncamento de primeira ordem. O estudo é limitado ao CIFAR-10 e perfis polinomiais, deixando a replicação em conjuntos de dados de maior resolução e a exploração de perfis aprendidos ou não polinomiais como trabalhos futuros.