Resumo Técnico: Da Composição de Especialistas Global para a Composição por Fatores em Modelos de Difusão Discreta

Declaração do Problema

Modelos de difusão discreta emergiram como um framework poderoso para tarefas de raciocínio complexas, particularmente através da geração composicional, que combina múltiplos modelos "especialistas" pré-treinados para generalizar além de suas distribuições de treinamento individuais. Avanços teóricos recentes (ex: SuperDiff, RNE, Corretores de Feynman-Kac) melhoraram a amostragem composicional ao introduzir pesos de mistura dependentes do tempo para alinhar a dinâmica composta com as distribuições alvo.

No entanto, esses métodos de estado da arte compartilham uma limitação fundamental: eles operam em uma base por amostra, atribuindo um único peso escalar global a cada especialista em cada passo temporal. Essa suposição "monolítica" trata implicitamente cada especialista como igualmente conhecedor em todo o espaço de estados. Em tarefas generativas complexas — como o raciocínio espacial, onde um especialista pode se especializar em bordas enquanto outro lida com texturas — esse peso global força um compromisso. Ele dilui o sinal do especialista correto em sua região específica de competência, falhando em capturar a independência espacial e funcional dos diferentes especialistas.

Metodologia: FactorDiff

Os autores propõem o FactorDiff, um framework de composição por fator que decompõe o processo de geração do nível da amostra para o nível do fator (por exemplo, pixels espaciais ou células de grade). Em vez de aplicar um peso único a todo um estado gerado, o FactorDiff roteia dinamicamente cada fator para o especialista mais relevante.

Componentes Principais

1. Roteamento por Fator: O espaço de estados é decomposto em fatores definidos pelo usuário. Durante o processo de amostragem, um mapa de roteamento atribui cada fator (ex: um pixel em uma grade) a um especialista específico. Isso permite contribuições heterogêneas de especialistas dentro de um único estado gerado.
2. Seleção Baseada em Confiança: O framework instancia o roteamento selecionando dinamicamente o especialista com maior confiança em cada posição durante a inferência. A variável de roteamento $z_v$ para uma variável $v$ é determinada por:
   $$z_v = \arg\max_{k} C(p_k(x_v | x_t))$$
   onde $C$ é uma métrica de confiança (a margem entre as duas principais probabilidades de predição). O passo reverso composto então seleciona a predição do especialista escolhido para aquela posição específica.
3. Arquitetura de Difusão Mascarada 2D Nativa: Para suportar isso, os autores introduzem uma arquitetura de difusão mascarada 2D que preserva a topologia da grade. Ao contrário de trabalhos anteriores que achatam grades em sequências 1D (introduzindo desalinhamentos topológicos), esta abordagem trata a grade como um grafo de rede (lattice graph).
   • Especialistas: Três tipos de modelos complementares são treinados: um Modelo Completo (todas as cores), um Modelo de Ocupação (fundo vs. primeiro plano) e um Modelo de Cor (apenas cores do primeiro plano).
   • Backbone: A rede de denoising utiliza uma arquitetura de refinamento iterativo (baseada em Convolução ou Transformer) com embeddings posicionais 2D aprendíveis para garantir a localidade espacial.

Justificativa Teórica

Métodos existentes frequentemente definem uma distribuição alvo como um produto ponderado de especialistas ($p(x) \propto \prod p_m(x)^{\alpha_m}$). Os autores argumentam que, para especialistas complementares, um expoente global $\alpha_m$ é mal especificado porque força as predições de um especialista (mesmo as errôneas fora de seu domínio) a influenciar todo o estado. O FactorDiff substitui isso por um produto por fator:
$$p_{comp}(x) \propto \prod_{f \in \mathcal{F}} \prod_{m=1}^M (p_m(x_f))^{\alpha_{m,f}}$$
Esta formulação permite que a expertise varie entre os fatores. No caso de roteamento rígido (hard-routing), $\alpha_{m,f}$ é um vetor one-hot, selecionando um único especialista por fator sem exigir normalização adicional ou correção de importância, já que o produto de condicionais válidos define uma transição conjunta válida.

Principais Contribuições

1. Framework FactorDiff: Um novo método de composição que trata a combinação de especialistas como um roteamento sobre fatores, especificamente instanciado com seleção de especialista ao nível de posição para modelos de difusão estruturados em grade.
2. Arquitetura de Difusão Mascarada 2D: Uma arquitetura nativa que mantém a topologia da grade, permitindo o treinamento de especialistas complementares (Ocupação, Cor, Completo) cujas predições podem ser compostas no tempo de inferência.
3. Validação Empírica no ARC-AGI: Avaliação extensiva no benchmark ARC-AGI demonstra que o roteamento simples por fator supera consistentemente esquemas de ponderação escalar global, particularmente em tarefas que exigem consistência lógica e desmembramento espacial.

Resultados Experimentais

Os autores avaliaram o FactorDiff em um subconjunto de 120 tarefas ARC-AGI usando o gerador RE-ARC para avaliação de distribuição interna (IID).

• Regime de Especialização (RQ1): Ao combinar especialistas altamente especializados (ex: Ocupação + Completo, ou Cor + Ocupação), o FactorDiff superou significativamente os baselines por amostra (SuperDiff, RNE, FKC).
  • Cor + Ocupação: Especialistas individuais pontuaram próximo de zero na tarefa completa (3,3% e 0,2% de acurácia exata). O FactorDiff os compôs para atingir 90,5% de acurácia exata, um enorme aumento sobre o melhor baseline por amostra (78,4%).
  • Ocupação + Completo: O FactorDiff alcançou 94,0% de acurácia exata, superando o melhor especialista individual (93,0%) e todos os métodos de composição por amostra.
• Regime de Generalista (RQ2): Ao combinar dois modelos generalistas fortes (Completo + Completo), o FactorDiff foi "seguro", igualando o melhor baseline por amostra (95,2%) sem degradar o desempenho. Isso indica que a seleção espacialmente variável não introduz instabilidade quando os especialistas concordam amplamente.
• Ablação: Mesmo ao compor dois modelos completos com objetivos idênticos, mas arquiteturas diferentes (Trans-IR vs. Conv-IR), o FactorDiff alcançou uma acurácia maior (94,8%) do que qualquer um dos modelos individuais, costurando com sucesso as regiões corretas de ambos.

Significância e Alegações

O artigo afirma que o FactorDiff fornece um mecanismo simples e eficaz para combinar especialistas arbitrários dentro de uma única trajetória de difusão. Sua principal significância reside em demonstrar que a ponderação global por amostra é insuficiente para tarefas onde os especialistas possuem domínios espacialmente ou funcionalmente especializados. Ao mover-se para a composição por fator, o método captura a expertise complementar que pesos globais não conseguem representar, alcançando uma generalização superior em tarefas que exigem rigorosa consistência lógica.

Os autores observam que, embora sua implementação atual dependa de roteamento baseado em confiança (assumindo que a confiança correlaciona-se com a correção), o framework se estende naturalmente para roteamento aprendido ou condicionado à tarefa. Eles posicionam este trabalho como um passo em direção a uma composição mais granular na difusão discreta, indo além das suposições "monolíticas" dos métodos atuais de estado da arte.