Calibri: Enhancing Diffusion Transformers via Parameter-Efficient Calibration

1. Problema e Motivação

Os Transformadores de Difusão (DiTs) tornaram-se o padrão de facto para geração de imagens (ex: Stable Diffusion 3, FLUX), substituindo arquiteturas baseadas em U-Net. No entanto, o artigo identifica que, apesar da arquitetura uniforme dos blocos DiT, suas contribuições funcionais são desiguais.

• Observação Chave: Estudos anteriores mostraram que certas camadas são "vitais", enquanto outras podem introduzir artefatos prejudiciais.
• Hipótese Central: A arquitetura DiT padrão possui um peso sub-ótimo entre seus blocos. O desempenho pode ser melhorado significativamente através de uma calibração post-hoc (após o treinamento) que ajuste a contribuição de cada bloco, sem a necessidade de re-treinar todo o modelo.
• Desafio: Métodos de alinhamento tradicionais (como RLHF, DPO) exigem fine-tuning completo, o que é computacionalmente caro e complexo.

2. Metodologia: Calibri

O Calibri é uma abordagem eficiente em parâmetros que trata a calibração como um problema de otimização de recompensa de caixa preta.

A. Formulação do Problema

O objetivo é encontrar um vetor de parâmetros de calibração $c^*$ que maximize uma função de recompensa $R(c)$, onde $R$ mede a qualidade da imagem gerada (usando modelos de recompensa como HPSv3, ImageReward, etc.).

• Parâmetros Otimizados: Apenas $\sim 10^2$ parâmetros (escalares) são ajustados, em vez de milhões.

B. Níveis de Granularidade

O método propõe três níveis de ajuste para os blocos internos do DiT:

1. Escalonamento de Bloco (Block Scaling): Ajusta uniformemente a saída de todas as camadas (Atenção e MLP) dentro de um único bloco arquitetural usando um único escalar.
2. Escalonamento de Camada (Layer Scaling): Ajusta cada camada individualmente dentro de um bloco com coeficientes distintos, oferecendo maior flexibilidade.
3. Escalonamento de Porta (Gate Scaling): Específico para arquiteturas multimodais (MM-DiT), onde tokens visuais e textuais possuem portas de modulação separadas ($\alpha, \beta, \gamma$), permitindo calibração especializada para cada modalidade.

C. Algoritmo de Otimização

Para resolver o problema de otimização sem gradiente, o Calibri utiliza a estratégia evolutiva CMA-ES (Covariance Matrix Adaptation Evolution Strategy).

• O algoritmo gera candidatos de parâmetros a partir de uma distribuição Gaussiana multivariada.
• Avalia os candidatos usando um modelo de recompensa.
• Atualiza iterativamente a média e a matriz de covariância para explorar o espaço de busca em direção às soluções ótimas.

D. Calibri Ensemble

Uma extensão do método que combina múltiplos modelos calibrados simultaneamente.

• Formulação: $F(x) = \sum \omega_i f_{\theta}^{s_i}(x)$.
• Permite agregar a diversidade de $N$ modelos otimizados, melhorando a robustez e a qualidade.
• Integra-se naturalmente com o paradigma de Classificador-Free Guidance (CFG), onde dois modelos (condicional e incondicional) são calibrados para otimizar a geração.

3. Resultados Experimentais

Os experimentos foram realizados em modelos de ponta: FLUX.1-dev, Stable Diffusion 3.5 Medium e Qwen-Image.

• Melhoria de Qualidade: O Calibri aumentou consistentemente as métricas de preferência humana (HPSv3), alinhamento de texto e qualidade de imagem (Q-Align, ImageReward) em todos os modelos base.
  • Exemplo (FLUX): HPSv3 subiu de 11.41 para 13.48.
• Redução de Passos de Inferência (NFE): Uma descoberta crucial é que o Calibri reduz drasticamente o número de passos necessários para gerar imagens de alta qualidade.
  • FLUX: De 30 passos para 15 passos.
  • SD-3.5M: De 80 passos para 30 passos.
  • Qwen-Image: De 100 passos para 30 passos.
• Comparação com Fine-tuning: Ao ser aplicado ao SD-3.5M, o Calibri alcançou desempenho comparável ao método Flow-GRPO (que ajusta 18,78 milhões de parâmetros), mas ajustando apenas 216 parâmetros.
• Validação Humana: Um estudo com 200 usuários mostrou preferência decisiva pelo Calibri em relação aos modelos originais tanto na preferência geral quanto no alinhamento com o texto.
• Custo Computacional: O custo de calibração é baixo (offline), variando de 32 a 356 horas de GPU H100, resultando em uma aceleração permanente de inferência de 2x a 3x.

4. Contribuições Principais

1. Descoberta de Sub-otimização: Evidência de que blocos DiT individuais podem ser reponderados (via escalonamento simples) para melhorar a geração, e que desabilitar certos blocos pode, paradoxalmente, melhorar a qualidade.
2. Método Calibri: Uma técnica de calibração post-hoc que otimiza apenas $\sim 10^2$ parâmetros via CMA-ES, evitando o custo de fine-tuning completo.
3. Eficiência de Inferência: Redução significativa no número de passos de difusão necessários, tornando a geração mais rápida e barata.
4. Versatilidade: Funciona em diferentes arquiteturas (DiT padrão e MM-DiT) e pode ser combinado com métodos de alinhamento existentes para ganhos adicionais.

5. Significado e Impacto

O trabalho demonstra que a arquitetura de modelos de difusão modernos possui um "potencial oculto" que não é totalmente explorado pelos pesos padrão. O Calibri oferece uma solução prática e econômica para elevar a qualidade de modelos de IA generativa, tornando-os mais rápidos e precisos sem a necessidade de re-treinamento massivo. Isso é particularmente relevante para aplicações do mundo real onde a latência e o custo computacional são críticos.