Resumo Técnico: Ensino de Distância Métrica a Modelos de Linguagem Autoregressivos Discretos

Declaração do Problema

Modelos de Linguagem de Grande Escala (LLMs) e outros modelos fundamentais autoregressivos discretos operam prevendo tokens de um vocabulário, tratando cada token como uma variável categórica discreta otimizada contra uma distribuição alvo one-hot. Embora essa formulação tenha permitido a adaptação de LLMs a diversos domínios, como visão, robótica e raciocínio multimodal, ela fundamentalmente ignora as relações métricas entre os tokens.

Em tarefas envolvendo elementos numericamente ou metricamente estruturados — como inteiros em raciocínio aritmético, coordenadas espaciais em detecção de objetos, ângulos de rotação em manipulação robótica ou embeddings quantizados em geração de imagens — a estrutura de distância intrínseca é perdida quando os tokens são reduzidos a alvos one-hot. O treinamento convencional ignora que uma previsão de "4" é semanticamente mais próxima de uma verdade fundamental de "5" do que uma previsão de "2". Essa limitação prejudica a eficiência de dados e o desempenho em domínios onde a distância é significativa.

Metodologia: DIST2Loss

Os autores introduzem o DIST2Loss (Função de Perda de Distância Discretizada), uma função objetivo consciente de distância para modelos autoregressivos discretos. A ideia central é substituir a distribuição alvo one-hot padrão por uma distribuição ponderada por recompensa derivada de distâncias de tokens predefinidas.

Formulação Central

Dado um subconjunto de vocabulário $V_d$ e um espaço métrico $(X, d)$, o DIST2Loss constrói uma distribuição alvo $p_d(v|x, t)$ para um token $v$ no passo de tempo $t$, dada uma token de verdade fundamental $x$. Esta distribuição é definida usando uma família exponencial discretizada:

$$p_d(v|x, t) = \frac{\exp\left(-\frac{d(v, x, t)}{\tau}\right)}{\sum_{v' \in V_d} \exp\left(-\frac{d(v', x, t)}{\tau}\right)}$$

Onde:

• $d(v, x, t)$ é a distância entre o token candidato $v$ e a verdade fundamental $x$ (por exemplo, distância Euclidiana para coordenadas, distância angular para rotações ou distância de embedding para códigos VQ).
• $\tau$ é um hiperparâmetro de temperatura que controla a suavidade da distribuição.

O objetivo final de treinamento combina a perda de entropia cruzada padrão ($L_{CE}$) com a regularização baseada em distância:
$$L = L_{CE} + \alpha L_{dist}$$
onde $L_{dist}$ é a divergência KL entre a distribuição prevista pelo modelo e o alvo consciente de distância $p_d$, e $\alpha$ equilibra os dois termos.

Interpretação Teórica

O DIST2Loss é interpretado como a solução de forma fechada para otimização de política regularizada por entropia. Em aprendizado por reforço (RL), a regularização por entropia encoraja a exploração, impedindo que a política colapse muito cedo. Sabe-se que a política ótima para maximizar a recompensa esperada com regularização por entropia é proporcional a $\exp(R(a)/\tau)$. O DIST2Loss instancia isso tratando tokens candidatos como ações e a distância negativa como um sinal de recompensa. Diferentemente do RL tradicional, que requer amostragem, simulações (rollouts) e atualizações iterativas, o DIST2Loss fornece uma distribuição alvo analítica, estável e sem variância, evitando a instabilidade associada aos estimadores de gradiente de política.

Lidando com Altas Dimensões e Sequências de Múltiplos Tokens

O framework é flexível o suficiente para lidar com vetores contínuos de alta dimensão (por exemplo, embeddings VQ) definindo métricas de distância sobre seus embeddings (por exemplo, similaridade cosseno). Para elementos estruturados de múltiplos tokens (como números de vários dígitos), os autores observam que aplicar a perda diretamente enfrenta um "problema de atribuição de crédito", onde o modelo não consegue distinguir qual token causou um erro. Para mitigar isso, eles propõem soluções práticas, como Aumento de Alvo Contrastivo (amostragem de vizinhos negativos) e Pesagem por Valor Posicional (ponderar a perda pela significância posicional), embora seus experimentos principais se concentrem em tarefas onde as recompensas se decompõem naturalmente no nível do token.

Principais Contribuições

1. Framework DIST2Loss: Um objetivo plug-and-play que integra estrutura métrica ao treinamento autoregressivo discreto sem exigir dados adicionais ou sobrecarga computacional significativa.
2. Ponte Teórica: Estabelecer uma conexão entre supervisão consciente de distância e RL regularizado por entropia, oferecendo uma alternativa estável ao aprendizado por reforço para tarefas com recompensas por token conhecidas.
3. Generalização: Demonstrar que o método é aplicável em diversos domínios, desde regressão escalar 1D até geração de imagens de alta dimensão, desde que os tokens admitam uma métrica significativa.

Resultados Experimentais

Os autores avaliaram o DIST2Loss em cinco domínios distintos, mostrando consistentemente melhorias na eficiência de dados e no desempenho downstream:

1. Regressão de Função Sintética (Tarefa de Brinquedo): Em um cenário de meta-aprendizado onde modelos aprendem regressão linear a partir de apenas 1–10 amostras, o DIST2Loss superou significativamente o SFT padrão e as linhas de base com vocabulário restrito em termos de Erro Absoluto Médio (MAE) e Raiz do Erro Quadrático Médio (RMSE), demonstrando generalização estrutural superior.
2. Ancoragem Visual: Ajuste fino do Phi3V em RefCOCO/RefCOCO+/RefCOCOg sem pré-treinamento em detecção de objetos, o DIST2Loss alcançou desempenho comparável a modelos state-of-the-art pré-treinados em conjuntos de dados de detecção em grande escala. Notavelmente, mesmo em "casos difíceis" onde ambos os modelos previram o objeto errado, as previsões do DIST2Loss estavam geometricamente mais próximas da verdade fundamental (maior IoU).
3. Manipulação Robótica: No VIMABench, o DIST2Loss melhorou as taxas de sucesso na geração de ações robóticas (ângulos de juntas e coordenadas). Os ganhos foram mais pronunciados em regimes com escassez de dados (1K amostras), destacando o valor de prios métricos quando os dados são limitados.
4. Modelagem de Recompensa Generativa: Aplicado ao alinhamento de LLM (RLHF), o DIST2Loss melhorou o desempenho de modelagem de recompensa no RewardBench e MT-Bench. A abordagem generativa com DIST2Loss superou tanto classificadores binários padrão quanto linhas de base SFT, particularmente nos domínios "Chat Hard" e "Safety".
5. Geração de Imagens: Usando representações quantizadas por vetores (VQ), o DIST2Loss aprimorou a aprendizagem de representações de imagens em modelos autoregressivos (LlamaGen), produzindo FID mais baixo e pontuações IS mais altas em comparação com linhas de base SFT em diferentes tamanhos de modelo e épocas de treinamento.

Significado e Alegações

O artigo afirma que supervisão consciente de distância oferece uma alternativa simples e geral à supervisão one-hot para modelos autoregressivos discretos. O significado reside na capacidade de:

• Aproveitar a Estrutura Intrínseca: Explorar as relações métricas inerentes a muitas tarefas não textuais (coordenadas, números, embeddings) que são de outra forma ignoradas pela entropia cruzada padrão.
• Melhorar a Eficiência de Dados: Alcançar melhor desempenho com menos amostras de treinamento codificando vieses indutivos diretamente na função de perda.
• Estabilizar o Treinamento: Fornecer uma alternativa de forma fechada e estável ao aprendizado por reforço para alinhamento de recompensas, evitando os problemas de amostragem e instabilidade dos gradientes de política.

Os autores enfatizam que o método é mais eficaz quando os tokens admitem uma métrica significativa (por exemplo, valores numéricos, coordenadas, embeddings quantizados). Em domínios sem geometria intrínseca, o método reduz-se naturalmente à supervisão one-hot padrão. O trabalho sugere que, para uma ampla gama de tarefas multimodais e de raciocínio, incorporar distância métrica é um componente crucial, mas anteriormente subutilizado, do treinamento de modelos fundamentais discretos.