Resumo Técnico: x-Prediction Is All You Need

Problema

Modelos de difusão e flow matching estabeleceram um desempenho de estado da arte em tarefas generativas, mas sofrem com altos custos de inferência, tipicamente exigindo dezenas a centenas de avaliações de função neural (NFEs) por amostra. Estratégias de aceleração existentes — como retificação de trajetória (Rectified Flow), destilação (progressiva, baseada em consistência, bootstrapping) e solvers de ODE avançados — geralmente requerem modificações na receita de treinamento, na arquitetura do modelo ou na própria trajetória. Consequentemente, esses métodos são frequentemente incompatíveis com a vasta biblioteca de checkpoints pré-treinados existentes (ex: SDXL, SD3.5M, FLUX) sem incorrer em custos de retreinamento. O campo carece de uma estratégia de aceleração livre de treinamento que funcione com os modelos "como eles são".

Metodologia: Truncated Jump Sampling (TJS)

O artigo propõe o Truncated Jump Sampling (TJS), uma estratégia de inferência livre de treinamento baseada no conceito de decodificabilidade de endpoint.

Insight Central: Decodificabilidade de Endpoint

Os autores formalizam que, para qualquer caminho de probabilidade afim não degenerado (a estrutura padrão para modelos de difusão e flow matching definida como $x_t = \alpha_t x_0 + \sigma_t \epsilon$), a amostra limpa $x_0$ pode ser recuperada unicamente de qualquer estado intermediário $x_t$ e sua velocidade de trajetória $u_t$ através de um decodificador de forma fechada.

• Recuperação Algébrica: Dado o sistema linear acoplado de posição e velocidade, $x_0$ é recuperado como $x_0 = (\sigma_t u_t - \dot{\sigma}_t x_t) / \Delta_t$, onde $\Delta_t$ é o determinante do caminho.
• Otimalidade: Sob objetivos de treinamento $\ell_2$ padrão, este decodificador algébrico corresponde ao estimador de Erro Quadrático Médio Mínimo (MMSE) $E[x_0 | x_t]$. Isso implica que modelos pré-treinados padrão aprendem implicitamente a prever o endpoint $x_0$ em cada timestep, mesmo que sejam treinados para prever ruído ou velocidade.
• Independência de Curvatura: Crucialmente, o erro desta previsão de endpoint depende apenas da precisão de estimativa do modelo e da incerteza irredutível do estado $x_t$, não da curvatura da trajetória ($\|\ddot{x}_t\|$). Isso desafia a suposição prevalecente de que trajetórias retas (como no Rectified Flow) são necessárias para uma geração precisa de poucos passos.

O Algoritmo

O TJS operacionaliza este insight interrompendo a integração da ODE precocemente em um tempo $t^*$ e imediatamente retornando o $x_0$ decodificado em vez de continuar até $t=0$.

1. Integração Parcial: O modelo integra a ODE para uma fração $\gamma$ dos passos totais ($k^* = \lceil \gamma K \rceil$).
2. Saída Antecipada (Early Exit): No passo $k^*$, o algoritmo computa a estimativa do endpoint $\hat{x}_0$ usando o campo de velocidade aprendido e o decodificador de forma fechada.
3. Terminação: O processo para, retornando $\hat{x}_0$ como a amostra final.

Esta abordagem requer zero retreinamento, zero destilação e zero mudanças de arquitetura. É compatível com qualquer solver de ODE, escala de classifier-free guidance (CFG) ou checkpoint existente.

Contribuições Principais

1. Formalização Teórica: O artigo estabelece a Decodificabilidade de Endpoint (Teorema 5) como uma propriedade universal de caminhos afins. Prova que o decodificador induzido é o estimador otimizado por MMSE (Teorema 6) e que a decomposição do erro é independente da curvatura da trajetória (Teorema 9).
2. Desafio à Retificação de Trajetória: Os autores provam que trajetórias retas são suficientes, mas não necessárias, para a previsão precisa do endpoint (Proposição de 12). Isso desacopla a aceleração de inferência da necessidade de retificação de trajetória, oferecendo uma base teórica de por que o TJS funciona em caminhos curvos onde a integração de Euler falharia.
3. Algoritmo Livre de Treinamento (TJS): Um amostrador prático que reduz as NFEs em 20–70% em diversas famílias de modelos sem modificar o modelo ou os dados de treinamento.
4. Validação Empírica Abrangente: Experimentos extensos em seis famílias de modelos (SDXL, SD3.5M, Z-Image-Turbo, ImageNet-256, CIFAR-10, MNIST) demonstram que o TJS alcança qualidade equivalente ou superior às trajetórias completas de ODE.

Resultados Experimentais

Os autores avaliaram o TJS em benchmarks de classe-condicional e texto-para-imagem:

• Redução de NFE: O TJS reduz as NFEs em 20–70% mantendo uma qualidade quase equivalente.
  • SDXL/SD3.5M: Alcança 95% da qualidade da ODE completa (medida por ImageReward) com aproximadamente 37–40% de economia de NFE (parando em $k^* \approx 18$ de 30 passos).
  • Z-Image-Turbo: Em um modelo destilado de 10 passos, o TJS alcança retenção de 95% de qualidade com 70% de economia adicional (parando em $k^*=2$, ou seja, 3 NFEs no total), demonstrando que a decodificabilidade de endpoint persiste mesmo em trajetórias destiladas.
  • Condicional de Classe: Em MNIST e CIFAR-10, o TJS às vezes supera a ODE completa (o fenômeno de "overshoot") porque evita os erros de discretização acumulados nos passos finais da integração de Euler.
• Hierarquia de Métricas: Os experimentos revelam uma hierarquia de convergência onde o conteúdo semântico (CLIP, PickScore) se resolve cedo (frequentemente por $k^* \approx 12$), enquanto detalhes estéticos refinados (ImageReward, AES) requerem integração mais profunda.
• Robustez: O desempenho do TJS é invariante à escolha do solver de ODE (DDIM, DPM++, UniPC, etc.), esquema de ruído (Beta, Karras, etc.) e escala de CFG, confirmando sua ortogonalidade aos componentes de inferência existentes.

Significância e Alegações

O artigo afirma que a aceleração moderada de poucos passos não requer retificação de trajetória, destilação ou redesenho de treinamento. A significância reside em reconhecer a decodificabilidade de endpoint como uma propriedade estrutural inerente aos modelos generativos em caminhos afins.

• Impacto Prático: O TJS fornece um caminho de aceleração imediato e de custo zero para a biblioteca existente de checkpoints ajustados (fine-tuned) da comunidade, que muitas vezes são inviáveis de retreinar ou destilar.
• Mudança Teórica: Ao provar que trajetórias retas não são necessárias para a previsão de endpoint, o trabalho desafia a motivação fundamental do Rectified Flow e dos Consistency Models em relação à aceleração de inferência.
• Natureza Complementar: Os autores posicionam o TJS como ortogonal à destilação. Enquanto a destilação comprime o comprimento total da trajetória, o TJS elimina a cauda redundante de qualquer trajetória restante, permitindo aceleração aditiva (ex: aplicar TJS a um modelo destilado como o Z-Image-Turbo).

O artigo conclui que reconhecer esta propriedade "não custa nada" e que ignorá-la deixa desempenho sobre a mesa, oferecendo uma maneira principiada de trocar profundidade de integração por velocidade de inferência sem sacrificar a qualidade da amostra gerada.