Resumo Técnico: MixLM

Declaração do Problema

Os Grandes Modelos de Linguagem (LLMs) demonstraram capacidades superiores de classificação de relevância em sistemas de busca e recomendação ao capturar sinais semânticos sutis que os métodos léxicos tradicionais perdem. No entanto, implantar LLMs como classificadores de cross-encoder em ambientes industriais enfrenta um gargalo crítico: o custo computacional.

As abordagens padrão de cross-encoder constroem prompts combinando consultas de usuários com descrições longas de itens candidatos (frequentemente milhares de tokens). Isso resulta em:

1. Prefill de Inferência Pesado: O custo computacional é dominado pela fase de prefill, que escala quadraticamente com o comprimento do contexto devido aos mecanismos de atenção.
2. Restrições de Throughput: Sob orçamentos estritos de latência (ex: 500ms), os sistemas são forçados a usar modelos menores e menos capazes ou a descartar características informativas (via sumarização de texto ou poda), sacrificando a qualidade da classificação.

O desafio central é reduzir significativamente o comprimento do contexto de entrada enquanto se preserva a rica interação semântica entre a consulta e o documento completo.

Metodologia: Estrutura MixLM

Os autores propõem o MixLM, uma estrutura de classificação baseada em LLM que substitui o texto longo do item por um conjunto compacto de tokens de embedding aprendidos. Esta arquitetura de "interação mista de texto-embedding" desacopla o processamento das descrições de itens da fase de classificação online.

1. Arquitetura

O MixLM consiste em dois componentes primários:

• Encoder LLM: Processa a descrição completa do item offline (ou nearline) para gerar uma representação compacta. Em vez de produzir um único vetor, ele produz uma sequência de tokens de embedding. Na produção, os autores utilizam last-$N$ sampling para selecionar um pequeno número de tokens (ex: $T_S=1$) dos estados ocultos do encoder para representar o item.
• Ranker LLM: Opera online. Ele recebe a consulta do usuário, texto auxiliar e os tokens de embedding pré-computados do encoder. Estes são concatenados para formar um prompt de "entrada mista", que é significativamente mais curto que o prompt de texto completo.

O ranker prevê a relevância ($p_{yes}$) com base nesta entrada mista. Como os embeddings do item são pré-computados e armazenados em cache, a inferência online processa apenas o prefixo da consulta e alguns tokens de embedding, reduzindo drasticamente o comprimento da sequência.

2. Pipeline de Treinamento

Os autores empregam uma receita de treinamento de três estágios para alinhar o encoder e o ranker:

• Estágio I (Fine-Tuning de Raciocínio de Domínio): Um LLM pré-treinado de 0.6B é ajustado em um conjunto de dados de raciocínio específico de domínio (180k amostras) usando traços de chain-of-thought de um juiz de relevância de 7B. Isso estabelece um modelo base capaz de compreender racionales de relevância.
• Estágio II (Treinamento do Professor de Classificação): O modelo do Estágio I é posteriormente ajustado como um classificador de texto puro usando descrições completas de itens (10.9M amostras). Este modelo serve como um "professor" de alta qualidade, mas é caro demais para o serviço online.
• Estágio III (Treinamento Conjunto de Encoder-Ranker): O encoder e o ranker são treinados conjuntamente usando a arquitetura de entrada mista. O objetivo de treinamento combina:
  • Perda de SFT de Soft-Label: Alinha as previsões com os rótulos de verdade fundamental (ground truth).
  • Perda de Destilação: Alinha a distribuição de saída do modelo de entrada mista com o modelo professor de texto puro (Estágio II).
  • Perda de Auto-Alinhamento: Garante que os outputs de embedding do encoder sejam compatíveis com os estados ocultos internos do ranker. Isso inclui alinhamento de estado oculto (combinando o último estado oculto da entrada mista com a entrada de texto completo) e alinhamento de predição.

Uma estratégia de aprendizado por currículo é usada no Estágio III, priorizando inicialmente o alinhamento de espaço (compatibilidade encoder-ranker) antes de mudar o foco para a destilação (precisão de predição de relevância).

3. Otimizações de Serventia

Para maximizar o throughput, os autores implementaram várias otimizações de nível de sistema:

• Shared-Prefix Prefill: Como o prefixo da consulta é idêntico para todos os itens em uma lista de classificação, o cache Key-Value (KV) para o prefixo é computado uma vez e reutilizado para todos os itens em um lote.
• In-Batch Prefix Caching: O mecanismo de inferência (SGLang) é modificado para fundir computações de atenção, permitindo que os tokens de sufixo (embeddings) atendam ao prefixo compartilhado em uma única passagem direta (forward pass).
• Arquitetura Multi-Processo: A pilha de serviço desacopla o frontend gRPC do motor de inferência usando multiprocessamento para contornar os gargalos do Global Interpreter Lock (GIL) do Python.
• Gerenciamento de Cache KV: Operações redundantes de cache são eliminadas para cargas de trabalho apenas de prefill para reduzir o overhead de CPU e aumentar a capacidade efetiva de memória GPU.

Principais Contribuições

1. Estrutura MixLM: Uma nova arquitetura de classificação de entrada mista que substitui documentos de texto longos por tokens de embedding compactos, reduzindo o comprimento do prompt em ordens de magnitude mantendo a expressividade semântica.
2. Pipeline de Treinamento de Ponta a Ponta: Uma estratégia de co-treinamento com perdas customizadas de destilação e auto-alinhamento que transfere com sucesso o conhecimento de um professor de texto completo de alta qualidade para um modelo estudante de entrada mista.
3. Stack de Serventia de Produção: Uma otimização abrangente do motor de inferência (incluindo cache de prefixo compartilhado e manipulação multi-processo) que permite o serviço de LLM com alto throughput.
4. Implantação Industrial: A primeira implantação de tráfego total de um sistema de classificação semântica baseado em LLM na busca de empregos do LinkedIn, validada através de testes A/B de larga escala.

Resultados Experimentais

Avaliação Offline

• Qualidade de Relevância: O MixLM alcançou um NDCG@10 de 0.9239, comparável ao modelo professor de texto puro (0.9432) e significativamente superior à recuperação baseada apenas em embeddings (0.8380) e às linhas de base de texto sumarizado (0.9218).
• Throughput: Sob um orçamento fixo de latência de 500ms:
  • O MixLM pontuou 22.000 itens/s/GPU.
  • Isso representa uma melhoria de 10.0× sobre a classificação de LLM com texto sumarizado (2.200 itens/s/GPU).
  • Isso representa uma melhoria de 75.9× sobre a classificação de LLM de texto completo (290 itens/s/GPU).

Teste A/B Online

• O MixLM foi implantado em 100% do tráfego para a busca de empregos impulsionada por IA do LinkedIn.
• O sistema manteve métricas de relevância em paridade com a linha de base de produção existente de escala limitada.
• O aumento do throughput permitiu a implantação total da classificação baseada em LLM, resultando em um aumento de 0,47% nos Usuários Ativos Diários (DAU).

Significância e Alegações

O artigo afirma que o MixLM fornece um caminho prático e escalável para implantar sistemas de classificação baseados em LLM em ambientes reais. Ao unir a riqueza semântica dos LLMs de texto completo com os requisitos de eficiência da busca industrial, o MixLM demonstra que:

1. Compressão de Contexto via Embeddings: É possível comprimir milhares de tokens em poucos tokens de embedding sem perda significativa na qualidade da classificação.
2. Co-Design de Sistema-Modelo: O alto throughput é alcançado não apenas por mudanças na arquitetura do modelo, mas através da integração estreita do design do modelo (entrada mista) com otimizações de infraestrutura (cache KV, batching).
3. Impacto de Negócio: Os ganhos de eficiência permitem a transição de experimentos de escala limitada ou de texto sumarizado para a implantação de LLM em tráfego total, impulsionando diretamente as métricas de engajamento do usuário (DAU) em uma rede profissional de grande escala.

Os autores posicionam este trabalho como uma solução para o trade-off "latência-throughput" que anteriormente impedia a adoção generalizada de cross-encoders de LLM em sistemas de busca de produção.