Published on

Escalando a Inferência de Modelos de Difusão Um Novo Paradigma

Autores
  • avatar
    Nome
    Ajax
    Twitter

Introdução

Avanços recentes em Grandes Modelos de Linguagem (LLMs) demonstraram a eficácia do escalonamento durante a inferência. Modelos como o1, o3, DeepSeek R1, QwQ e Step Reasoner mini mostraram que o aumento da computação durante a inferência pode melhorar significativamente o desempenho. Isso levanta a questão: esse princípio também pode ser aplicado a modelos de difusão?

Uma equipe liderada por Xie Saining na Universidade de Nova York explorou essa questão. Através de uma investigação sistemática usando uma estrutura de busca geral, eles descobriram que o escalonamento no tempo de inferência é realmente eficaz para modelos de difusão. Aumentar a computação durante a inferência leva a uma melhoria significativa na qualidade das amostras geradas. Além disso, a complexidade das imagens permite diferentes combinações de componentes dentro da estrutura, adaptadas a vários cenários de aplicação.

Principais Descobertas

  • O escalonamento no tempo de inferência é eficaz para modelos de difusão: Alocar mais recursos computacionais durante a inferência resulta em amostras de maior qualidade.
  • Flexibilidade em combinações de componentes: A estrutura permite diferentes configurações de componentes, atendendo a várias aplicações.
  • Além das Etapas de Denoising: A pesquisa sugere que buscar melhor ruído durante a amostragem é outra dimensão para escalar o NFE, além de simplesmente aumentar as etapas de denoising.
  • Dois Eixos de Design: A estrutura se concentra em dois eixos de design principais:
    • Verificadores: Fornecer feedback durante o processo de busca.
    • Algoritmos: Encontrar melhores candidatos a ruído.

Metodologia de Pesquisa

A equipe explorou três cenários diferentes para verificadores, simulando vários casos de uso:

  1. Cenários onde informações privilegiadas sobre a avaliação final estão disponíveis.
  2. Cenários onde informações condicionais estão disponíveis para guiar a geração.
  3. Cenários sem informações adicionais disponíveis.

Para algoritmos, eles investigaram:

  1. Busca Aleatória: Selecionar o melhor de um conjunto fixo de candidatos.
  2. Busca de Ordem Zero: Melhorar iterativamente os candidatos a ruído usando feedback do verificador.
  3. Busca de Caminho: Melhorar iterativamente as trajetórias de amostragem de difusão usando feedback do verificador.

Inicialmente, o estudo explorou esses designs em uma configuração relativamente simples de geração condicional de classe ImageNet. Posteriormente, eles aplicaram esses designs à geração condicional de texto em maior escala e avaliaram sua estrutura proposta.

Escalonando o Tempo de Inferência

O artigo propõe uma estrutura para escalar o tempo de inferência em modelos de difusão, enquadrando o desafio como uma busca por ruído de amostragem ideal. O processo envolve dois componentes principais:

  • Verificadores: São modelos pré-treinados que avaliam a qualidade das amostras geradas. Eles recebem amostras geradas, juntamente com condições opcionais, e geram uma pontuação escalar.
  • Algoritmos: Esses algoritmos usam as pontuações do verificador para encontrar melhores amostras candidatas. A função recebe um verificador (V), um modelo de difusão pré-treinado (_θ) e um conjunto de amostras e condições geradas, produzindo o melhor ruído inicial.

O orçamento total de inferência é medido pelo número total de avaliações de função (NFE), incluindo etapas de denoising e custos de busca.

Verificadores de Busca

Os pesquisadores começaram com um verificador Oracle, que tem informações completas sobre a avaliação final das amostras selecionadas. Para ImageNet, isso incluía métricas como FID e IS. Em seguida, eles exploraram modelos pré-treinados mais acessíveis como verificadores supervisionados, como CLIP e DINO. Esses modelos foram usados para classificar amostras, selecionando a amostra com o logit mais alto correspondente ao rótulo da classe.

No entanto, eles observaram que esses classificadores, operando ponto a ponto, apenas parcialmente se alinham com os objetivos da pontuação FID. Isso levou a uma redução na variação da amostra e ao colapso do modo à medida que a computação aumentava. Esse fenômeno, denominado "hacking do verificador", foi acelerado pelo espaço de busca irrestrito do algoritmo de busca aleatória.

Curiosamente, o estudo descobriu que os verificadores não precisam necessariamente de informações condicionais para orientar efetivamente a busca. Eles observaram uma forte correlação entre os logits dos classificadores DINO/CLIP e a similaridade de cosseno do espaço de recursos entre a previsão x em um nível de baixo ruído e a amostra limpa final. Isso levou ao uso de verificadores auto-supervisionados, que não exigem informações condicionais adicionais e ainda exibem um comportamento de escalonamento eficaz.

Algoritmos de Busca

Para mitigar o hacking do verificador, os pesquisadores exploraram algoritmos de busca mais refinados que otimizam gradualmente as amostras candidatas. Isso incluiu um método de busca de ordem zero:

  1. Começar com um ruído Gaussiano aleatório como um ponto pivô.
  2. Encontrar N candidatos na vizinhança do ponto pivô.
  3. Executar os candidatos através do solucionador ODE para obter amostras e pontuações do verificador.
  4. Atualizar o ponto pivô com o melhor candidato e repetir as etapas 1-3.

Eles também investigaram um algoritmo de busca de caminho, que explora a possibilidade de buscar ao longo da trajetória de amostragem:

  1. Amostrar N amostras de ruído inicial e executar o solucionador ODE até um nível de ruído σ.
  2. Adicionar ruído a cada amostra e simular um processo de ruído para frente.
  3. Executar um solucionador ODE em cada amostra ruidosa e manter os N melhores candidatos com base nas pontuações do verificador, repetindo até que o solucionador ODE atinja σ=0.
  4. Buscar aleatoriamente as N amostras restantes e manter a melhor.

Tanto os algoritmos de busca de ordem zero quanto os de busca de caminho mantêm uma forte localidade em comparação com a busca aleatória.

Escalonamento em Cenários de Texto para Imagem

A equipe examinou as capacidades de escalonamento da estrutura de busca em tarefas de texto para imagem em maior escala. Eles usaram os conjuntos de dados DrawBench e T2I-CompBench para avaliação, com o modelo FLUX.1-dev como base. Eles também expandiram a seleção de verificadores supervisionados, incluindo Aesthetic Score Predictor, CLIPScore e ImageReward. Além disso, eles criaram um Verifier Ensemble combinando esses três verificadores.

Análise: Alinhamento Verificador-Tarefa

O estudo comparou os resultados de várias combinações de verificador-algoritmo em diferentes conjuntos de dados. No DrawBench, eles descobriram que usar todos os verificadores geralmente melhorava a qualidade da amostra. No entanto, eles observaram que usar os verificadores Aesthetic e CLIP isoladamente poderia levar à sobreajuste de seus vieses, resultando em impactos negativos um sobre o outro. Isso decorre de uma incompatibilidade em seu foco de avaliação: o Aesthetic Score se concentra na qualidade visual, muitas vezes favorecendo imagens altamente estilizadas, enquanto o CLIP prioriza o alinhamento visual-texto, às vezes sacrificando a qualidade visual. Eles notaram que alguns verificadores são mais adequados para tarefas específicas, e a eficácia de um verificador depende de seu alinhamento com os requisitos da tarefa.

Desempenho do Algoritmo

Os três algoritmos de busca (Aleatório, Ordem Zero e Caminho) melhoraram efetivamente a qualidade da amostragem no DrawBench. No entanto, a Busca Aleatória superou em alguns aspectos devido à natureza local dos outros dois métodos. A busca aleatória convergiu mais rapidamente para o viés do verificador, enquanto os outros dois algoritmos exigem melhoria em candidatos menos do que ideais.

Compatibilidade com Ajuste Fino

A equipe investigou a compatibilidade de seu método de busca com modelos ajustados. Eles usaram um modelo Stable Diffusion XL ajustado com DPO e descobriram que o método de busca poderia ser generalizado para diferentes modelos e melhorar o desempenho de modelos já alinhados.

Efeitos de Diferentes Dimensões da Computação de Inferência

O estudo explorou como diferentes aspectos da computação de inferência afetam os resultados:

  • Número de iterações de busca: Aumentar as iterações aproxima o ruído do ótimo.
  • Computação por iteração de busca: Ajustar o número de etapas de denoising por iteração revela diferentes regiões computacionalmente ótimas.
  • Computação de geração final: A equipe usou configurações ideais para as etapas finais de denoising para garantir a mais alta qualidade final da amostra.

Eficácia do Investimento em Computação

Os pesquisadores exploraram a eficácia do escalonamento no tempo de inferência em modelos de difusão menores. Eles descobriram que, para ImageNet, o escalonamento de modelos menores pode ser muito eficiente. Em certos casos, buscar em um modelo menor pode superar modelos maiores sem busca. No entanto, a eficácia depende do desempenho de linha de base do modelo menor.

Em configurações baseadas em texto, o PixArt-Σ, usando apenas uma fração da computação, superou o FLUX-1.dev. Esses resultados demonstram que recursos computacionais significativos gastos durante o treinamento podem ser compensados por pequenas quantidades de computação durante a geração, resultando em amostras de maior qualidade de forma mais eficiente.