Resumo Técnico: Aprendizado Multi-Objetivo para Modelos de Difusão sob Aprendizado Semi-Supervisionado

1. Formulação do Problema

Os modelos de difusão estão sendo cada vez mais implantados como geradores condicionais para diversas tarefas, como síntese de texto para imagem em domínios de prompts heterogêneos ou controle robótico em múltiplos ambientes. Esses cenários constituem naturalmente um problema de Aprendizado Multi-Objetivo (MOL), onde um único modelo deve alcançar desempenho Pareto-ótimo em múltiplas distribuições-alvo.

Um desafio central nesse cenário é o custo estatístico de treinar um modelo "generalista". Para lidar com múltiplos objetivos, a classe de modelos generalistas ($\mathcal{F}$) tipicamente requer capacidade substancialmente maior do que as classes de modelos "especialistas" ($\mathcal{H}_k$) necessárias para tarefas individuais. Em cenários totalmente supervisionados, a complexidade de amostragem para treinar o generalista escala com sua própria complexidade ($C_\mathcal{F}$), o que frequentemente é proibitivamente alto.

Além disso, implantações no mundo real frequentemente enfrentam um regime de dados assimétrico:

• Dados pareados $(x, y)$ (por exemplo, pares texto-imagem curados ou demonstrações de robôs especialistas) são escassos e caros de adquirir.
• Dados de condição não pareados $y$ (por exemplo, prompts de texto brutos ou observações de estado) são abundantes.

O artigo aborda a seguinte questão: Podemos aprender um modelo de difusão Pareto-ótimo de maneira eficiente em termos de amostras, aproveitando dados fracamente supervisionados (condição não rotulada)?

2. Metodologia

Os autores propõem um framework de treinamento semi-supervisionado em duas etapas que reconcilia a necessidade de um generalista de alta capacidade com a escassez de dados pareados.

Etapa 1: Treinamento de Especialistas

Para cada um dos $K$ objetivos (distribuições-alvo $P_k$), uma função de pontuação especialista leve $h_k \in \mathcal{H}_k$ é treinada utilizando apenas os dados pareados limitados disponíveis $\{(x_i^k, y_i^k)\}_{i=1}^n$.

• Assume-se que as classes de especialistas $\mathcal{H}_k$ tenham complexidade significativamente menor ($C_{\mathcal{H}_k} \ll C_\mathcal{F}$) do que a classe generalista.
• O treinamento minimiza o objetivo de correspondência de pontuação de remoção de ruído (denoising score matching) nos dados pareados.

Etapa 2: Treinamento do Generalista via Amostragem Pseudo

Um modelo generalista grande $f \in \mathcal{F}$ é treinado para aproximar o conjunto Pareto-ótimo em todos os objetivos.

• Geração de Pseudo-Amostras: Os especialistas treinados $\{h_k\}$ são usados para gerar dados pseudo-pareados. Para cada condição não rotulada abundante $y_i^k \sim P_Y^k$, o especialista gera um pseudo-alvo $\tilde{x}_i^k \sim P_{h_k}(\cdot|y_i^k)$.
• Objetivo: O generalista é treinado em uma combinação dos dados rotulados originais e dos dados pseudo-pareados gerados. O objetivo de treinamento minimiza uma escalarização $S$ das perdas de correspondência de pontuação em todos os $K$ objetivos:
  $$\hat{f}_S = \arg\min_{f \in \mathcal{F}} S(\tilde{L}_1(f), \dots, \tilde{L}_K(f))$$
  onde $\tilde{L}_k$ é a perda empírica calculada sobre as pseudo-amostras (com um termo de correção para contabilizar o viés do especialista).

Extensão para Tomada de Decisão Sequencial

O framework é estendido para políticas de difusão em aprendizado por imitação. Aqui, a "condição" é o estado $s$, e o "alvo" é a ação $a$. Os autores abordam a mudança de distribuição inerente aos rolagens (rollouts) on-policy: como as distribuições de estado de especialistas são difíceis de amostrar diretamente, a segunda etapa usa trajetórias geradas pelas políticas especialistas aprendidas sob as dinâmicas do ambiente para coletar dados pseudo para a política generalista.

3. Principais Contribuições Teóricas

Limites de Complexidade de Amostragem

O artigo estabelece limites de generalização não assintóticos para a abordagem semi-supervisionada proposta.

• Resultado Principal (Teorema 3.2): O número de amostras pareadas (rotuladas) necessárias para alcançar um erro alvo $\epsilon$ depende apenas da complexidade das classes de especialistas ($C_{\mathcal{H}_k}$), e não da classe generalista ($C_\mathcal{F}$).
  • Amostras rotuladas necessárias: $\tilde{O}\left(\frac{K C_\mathcal{H}}{\epsilon^2}\right)$.
  • Condições não rotuladas necessárias: $\tilde{O}\left(\frac{K C_\mathcal{F}}{\epsilon^2}\right)$.
• Comparação: Em contraste, o treinamento totalmente supervisionado padrão do generalista exigiria $\tilde{O}\left(\frac{K C_\mathcal{F}}{\epsilon^2}\right)$ amostras rotuladas. Isso demonstra uma melhoria significativa na eficiência de amostragem quando $C_\mathcal{F} \gg C_\mathcal{H}$.
• Escalarizações Lineares (Teorema 3.5): Para escalarizações lineares (soma ponderada) e classes de funções convexas, os autores derivam taxas mais precisas, refinando ainda mais o limite sobre o erro de aproximação.

Estimativa de Distribuição e Garantias de Política

• Estimativa de Distribuição: O artigo preenche a lacuna entre o erro de correspondência de pontuação e a estimativa de distribuição (distância de Variação Total). Ele prova que o generalista aprendido alcança desempenho Pareto-ótimo na aproximação das distribuições condicionais alvo.
• Políticas de Difusão (Teorema 3.6): Os autores fornecem a primeira garantia teórica sobre a lacuna de sub-otimalidade de políticas de difusão em aprendizado por imitação sob cenários semi-supervisionados. Eles mostram que o framework lida com sucesso com a mudança de distribuição causada por rolagens on-policy, mantendo ao mesmo tempo os benefícios de eficiência de amostragem do pipeline de especialista para generalista.

4. Resultados Experimentais

Os autores validam sua teoria em dois domínios distintos:

1. Manipulação Robótica (StackCube & PickCube):

   • Configuração: Treinamento de políticas de difusão para tarefas sob randomização de domínio variável (cenário, iluminação, pose da câmera).
   • Resultado: O método MOL semi-supervisionado supera significativamente uma linha de base apenas com rótulos, particularmente em cenários Fora de Distribuição (OOD) (por exemplo, +16% na taxa de sucesso no StackCube Nível 2, +14% no Nível 3). Isso sugere que as pseudo-rolagens fornecem supervisão crucial para lidar com mudanças de distribuição.

2. Restauração de Imagem (Inpainting CelebA-HQ):

   • Configuração: Preenchimento de rostos com três tipos diferentes de máscaras (linhas alternadas, região do rosto, traço largo).
   • Resultado: A abordagem semi-supervisionada melhorou o PSNR, SSIM e LPIPS em todos os tipos de máscara em comparação com a linha de base apenas com rótulos, demonstrando a eficácia do método em tarefas generativas de imagem.

5. Significado e Alegações

O artigo afirma fornecer um framework estatístico principiado para aprendizado multi-objetivo em modelos de difusão sob regimes semi-supervisionados. Seu significado principal reside em:

• Ponte entre Teoria e Prática: Formaliza o pipeline de destilação "de especialista para generalista", oferecendo uma justificativa teórica para por que essa abordagem é mais eficiente em amostras do que o aprendizado multi-tarefa direto quando os dados rotulados são escassos.
• Eficiência de Amostragem: Prova que o custo estatístico de aprender um generalista complexo pode ser desacoplado de sua própria capacidade, dependendo em vez disso da complexidade de especialistas leves treinados em dados limitados.
• Lidar com Mudança de Distribuição: Estende essas garantias para tomada de decisão sequencial (aprendizado por imitação), abordando um desafio prático crítico em robótica.

Os autores mantêm uma postura modesta, reconhecendo que suas garantias dependem de uma hipótese de realizabilidade (de que as funções de pontuação verdadeiras residem dentro das classes de especialistas) e não abordam atualmente a generalização Fora de Distribuição (OOD) além dos níveis específicos de randomização de domínio testados. Eles veem este trabalho como um passo fundamental para entender as propriedades estatísticas do MOL semi-supervisionado para modelos generativos.