Resumo Técnico: STEP-OPD

Declaração do Problema

A destilação on-policy (OPD) emergiu como uma abordagem viável para consolidar múltiplos modelos de geração de imagem especializados em tarefas em um único modelo estudante, supervisionando as próprias trajetórias de denoising do estudante. No entanto, os métodos de OPD existentes sofrem de duas limitações primárias:

1. Teto de Correspondência com o Professor (Teacher-Matching Ceiling): Os métodos atuais otimizam o estudante unicamente para corresponder à velocidade de saída do professor nos estados visitados. Isso trata o professor como o limite superior absoluto do objetivo de otimização. Como ilustrado no artigo, a perda padrão de correspondência de velocidade ao quadrado é minimizada exatamente quando a previsão do estudante é igual à previsão do professor; qualquer tentativa de "ir além" do professor aumenta a perda, impedindo que o estudante alcance capacidades superiores às do professor.
2. Dinâmica Interna Sub-Determinada (Under-Constrained Internal Dynamics): A supervisão de velocidade ao nível da saída fornece apenas um sinal de ponto final. Como a velocidade final agrega os efeitos de todos os blocos da rede, desvios nas atualizações de representação entre diferentes camadas podem se cancelar mutuamente na previsão final. Consequentemente, um estudante pode alcançar uma baixa perda de saída através de compensação sem aprender fielmente o processamento estruturado de tokens visuais e informações de condicionamento do professor. Isso enfraquece a transferência de capacidades que exigem desenvolvimento progressivo através das camadas, como vinculação de atributos (attribute binding), composição espacial e renderização de texto.

Metodologia: STEP-OPD

Os autores propõem o STEP-OPD (Step-wise Output Extrapolation and Representation Change Alignment), um framework que estende o alvo de aprendizado do estudante além do professor e introduz restrições explícitas sobre a evolução da representação interna. O método consiste em dois componentes principais:

1. Extrapolação de Saída Específica de Tarefa (Task-Specific Output Extrapolation)

Em vez de usar a velocidade do professor ($v_{T_k}$) como o alvo final, o STEP-OPD utiliza a diferença de velocidade entre o professor específico da tarefa e o modelo base compartilhado ($v_B$) como uma direção para aprendizado adicional.

• Mecanismo: A velocidade de alvo extrapolada $\hat{v}_{k,i}$ é construída adicionando uma versão escalonada da diferença professor-base à velocidade do professor:
  $$\hat{v}_{k,i} = v_{T_k,i} + \alpha_k (v_{T_k,i} - v_{B,i})$$
  onde $\alpha_k \geq 0$ é um coeficiente de extrapolação específico da tarefa.
• Efeito: Quando $\alpha_k > 0$, o alvo de aprendizado desloca-se além do professor ao longo da direção de melhoria introduzida pelo pós-treinamento. Isso transforma o professor de um ponto final fixo em um guia direcional para otimização adicional.
• Estratégia de Treinamento: Para evitar instabilidade, o coeficiente $\alpha_k$ é ativado gradualmente de 0 a um valor máximo durante um período de warm-up.

2. Alinhamento de Mudança de Representação Intermediária (RCA)

Para abordar a dinâmica interna sub-determinada, o STEP-OPD alinha a evolução dos estados ocultos entre o estudante e o professor através dos blocos da rede.

• Mecanismo: Em vez de corresponder aos estados ocultos absolutos, o método alinha a direção e a magnitude das mudanças nas representações ocultas entre blocos Transformer adjacentes. Para um par de blocos $p=(l_0, l_1)$, a mudança é definida como $\Delta h^p_{M,i} = h^{l_1}_{M,i} - h^{l_0}_{M,i}$.
• Função de Perda: O alinhamento de perda consiste em duas partes:
  1. Perda de Direção: Minimiza o MSE entre as mudanças normalizadas por RMS por token do estudante e do professor.
  2. Perda de Magnitude: Minimiza a diferença ao quadrado entre as magnitudes RMS em escala logarítmica das mudanças.
• Integração: Esta supervisão interna é aplicada diretamente às mudanças do professor (não extrapoladas), garantindo que o estudante aprenda a dinâmica de transformação específica do professor, enquanto o alvo de saída é permitido exceder o professor.

O objetivo final de treinamento combina a perda de saída extrapolada ($L_{out}$) e a perda de mudança de estado oculto ($L_{hidden}$):
$$L = L_{out} + \lambda_h L_{hidden}$$

Principais Contribuições

1. Proposta de Framework: Introdução do STEP-OPD, um framework de destilação on-policy que permite o aprendizado "além do professor" enquanto transfere a dinâmica de representação interna.
2. Extrapolação de Saída: Uma técnica inovadora que constrói alvos de aprendizado além do professor, aproveitando a diferença de velocidade entre o modelo base e os professores específicos de cada tarefa.
3. Alinhamento de Mudança de Representação: Um método para supervisionar o cálculo interno do estudante ao corresponder à direção e magnitude das atualizações de representação entre camadas, abordando as limitações da supervisão apenas de saída.
4. Validação Empírica: Experimentos abrangentes demonstrando que o estudante unificado STEP-OPD supera tanto os baselines de OPD padrão quanto os respectivos professores de tarefa única.

Resultados Experimentais

Experimentos foram conduzidos no Stable Diffusion 3.5 Medium em três grupos de capacidade: alinhamento composicional (GenEval), renderização de texto (OCR) e preferência humana (PickScore, Aesthetic, HPSv2.1, ImageReward).

• Ganhos de Desempenho: Quando aplicado ao DiffusionOPD, o STEP-OPD aumentou a pontuação geral do GenEval de 0.927 para 0.961 e a pontuação OCR de 0.941 para 0.946. Também melhorou todas as métricas de preferência humana.
• Superando os Professores: O estudante unificado STEP-OPD superou os respectivos professores de tarefa única em todas as três categorias de capacidade (ex: 0.961 vs. 0.935 no GenEval), demonstrando consolidação de capacidade eficaz e desempenho além do professor.
• Estudos de Ablação:
  • Alinhamento de Mudança de Representação (RCA): Sozinho, o RCA melhorou o GenEval (0.927 $\to$ 0.959) e o OCR (0.941 $\to$ 0.944), confirmando que a supervisão interna é crucial para tarefas composicionais e de texto.
  • Extrapolação de Saída (OE): Sozinho, o OE proporcionou ganhos significativos em métricas baseadas em preferência e melhorou ainda mais o GenEval/OCR, validando a hipótesia "além do professor".
  • Combinado: A combinação de ambos os componentes produziu os resultados mais fortes.
• Análise Qualitativa: Comparações visuais mostraram que o STEP-OPD produziu adesão ao prompt mais precisa, melhor renderização de texto e maior qualidade visual em comparação com os baselines. Especificamente, o RCA reduziu estruturas internas espúrias, enquanto o OE melhorou detalhes finos e texturas de materiais.

Significância e Alegações

O artigo afirma que o STEP-OPD altera fundamentalmente o paradigma da destilação on-policy de correspondência com o professor para extrapolação guiada pelo professor. Ao desacoplar o alvo de saída da velocidade exata do professor e supervisionar explicitamente a dinâmica de transformação interna, o método permite que um único modelo estudante consolide múltiplas capacidades especializadas e exceda o desempenho dos professores individuais que geraram os dados de treinamento. Os autores afirmam que esta abordagem resolve as limitações da supervisão apenas de saída e fornece um mecanismo robusto para transferir capacidades generativas complexas e dependentes de camadas.