💬 NLP

Vectorizing the Trie: Efficient Constrained Decoding for LLM-based Generative Retrieval on Accelerators

Este trabalho apresenta o STATIC, uma técnica de decodificação restrita altamente eficiente para aceleradores de hardware que transforma árvores de prefixo em matrizes esparsas vetorializadas, permitindo a primeira implantação em escala industrial de recuperação generativa estritamente restrita com impacto mínimo na latência.

Autores originais: Zhengyang Su, Isay Katsman, Yueqi Wang, Ruining He, Lukasz Heldt, Raghunandan Keshavan, Shao-Chuan Wang, Xinyang Yi, Mingyan Gao, Onkar Dalal, Lichan Hong, Ed Chi, Ningren Han

Publicado 2026-02-27

📖 4 min de leitura☕ Leitura rápida

CC BY 4.0

Autores originais: Zhengyang Su, Isay Katsman, Yueqi Wang, Ruining He, Lukasz Heldt, Raghunandan Keshavan, Shao-Chuan Wang, Xinyang Yi, Mingyan Gao, Onkar Dalal, Lichan Hong, Ed Chi, Ningren Han

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você tem um chef de cozinha super inteligente (o Modelo de Linguagem ou LLM) que trabalha no maior restaurante do mundo (o YouTube). A tarefa dele é criar pratos (recomendar vídeos) para milhões de clientes.

No passado, esse chef usava um sistema antigo: ele olhava para uma lista gigante de ingredientes e tentava adivinhar qual prato combinaria melhor com o cliente. Mas agora, o chef aprendeu a criar o prato do zero, escrevendo o nome do prato palavra por palavra, como se estivesse ditando uma receita. Isso é chamado de "Recuperação Generativa".

O problema é que, às vezes, o chef é tão criativo que começa a inventar pratos que não existem na cozinha, ou que o cliente não pode comer (por exemplo, um prato que acabou de ser preparado e ainda está quente demais, ou um ingrediente que está fora de estoque).

O Problema: O Labirinto de Papel

Para evitar que o chef invente coisas ruins, os donos do restaurante dizem: "Só podemos servir pratos que tenham ingredientes frescos dos últimos 7 dias".

A maneira tradicional de fazer isso era dar ao chef um mapa de labirinto de papel (chamado de Trie ou Árvore de Prefixo). A cada palavra que o chef escrevia, ele tinha que pegar o mapa, procurar a seta correta, andar até a próxima sala e ver quais eram as opções permitidas.

Onde estava o problema?
Esse mapa de papel era ótimo para humanos, mas péssimo para os robôs super-rápidos (os aceleradores de hardware como TPUs e GPUs) que cozinham os pratos.

Caminhada desajeitada: O robô tinha que andar de um lado para o outro no mapa, pulando de um papel para outro. Isso fazia o robô perder muito tempo apenas "andando" e não "cozinhando".
Memória bagunçada: O robô precisava ir buscar informações em lugares aleatórios da memória, o que deixava o sistema lento e travado.

A Solução: O Menu Digital Vetorizado (STATIC)

A equipe do Google criou uma solução chamada STATIC. Em vez de dar ao robô um labirinto de papel para ele andar, eles transformaram todo o mapa em um menu digital inteligente e plano.

Aqui está como funciona a analogia:

Do Labirinto para a Planilha: Em vez de ter que seguir setas e virar esquinas (o que é lento), eles transformaram o mapa inteiro em uma grande planilha de dados (uma matriz esparsa).
O Pulo do Gato (Operações Vetorizadas): Imagine que, em vez de o robô andar até a próxima opção, ele simplesmente olha para a planilha e, de uma só vez, vê todas as opções permitidas para o próximo passo. É como se o robô tivesse visão de raio-X e pudesse ler 100 opções ao mesmo tempo, em vez de ler uma por uma.
Sem Trânsito: Como tudo está organizado em blocos contíguos (juntos), o robô não precisa "viajar" pela memória. Ele puxa os dados de uma vez só, como se estivesse bebendo de um cano largo em vez de um cano fino.

Por que isso é incrível?

Velocidade Relâmpago: O método antigo (no processador comum) era 1.000 vezes mais lento. O novo método é tão rápido que o tempo extra que ele gasta é quase imperceptível (menos de 0,03 milissegundos por passo). É como trocar de andar a pé para usar um foguete.
Frescor Garantido: Agora, o chef só consegue escrever nomes de pratos que realmente existem e que estão frescos. Se um vídeo foi postado há 10 dias, o sistema bloqueia a palavra "fresco" antes mesmo do chef pensar em usá-la.
Escala Gigantesca: Isso funciona mesmo quando o restaurante tem 20 milhões de opções de pratos. O sistema não fica lento, ele continua voando.

O Resultado no Mundo Real

No YouTube, eles testaram isso. Antes, o sistema às vezes recomendava vídeos velhos ou fora de estoque, frustrando os usuários. Com o STATIC:

O sistema passou a recomendar 5% mais vídeos novos (dos últimos 7 dias).
As pessoas clicaram mais nos vídeos (o que significa que estavam mais felizes).
Tudo isso aconteceu sem deixar o aplicativo lento para o usuário.

Resumo da Ópera

O papel do STATIC foi pegar uma ideia complexa de "navegar em árvores de dados" (que era lenta e bagunçada para computadores modernos) e transformá-la em uma operação matemática simples e super rápida (como ler uma lista).

É como se, em vez de pedir para um entregador correr por uma cidade inteira para pegar um pacote, você tivesse um drone que puxasse todos os pacotes de uma vez só de um armazém organizado. O resultado? Mais velocidade, menos erros e clientes mais felizes.

Título: Vetorização do Trie: Decodificação Constrained Eficiente para Recuperação Generativa Baseada em LLM em Aceleradores

1. O Problema

A Recuperação Generativa (Generative Retrieval - GR) emergiu como um paradigma poderoso para sistemas de recomendação baseados em Grandes Modelos de Linguagem (LLMs). Nesses modelos, os itens são representados por sequências de tokens discretos (IDs Semânticos ou SIDs), e o modelo gera esses tokens autoregressivamente.

No entanto, existem dois desafios críticos para a adoção industrial dessa tecnologia:

Controle do Espaço de Saída: Sistemas de recomendação reais exigem restrições de negócios estritas (ex.: apenas itens "em estoque", vídeos "recentes" (freshness), ou categorias específicas). Modelos generativos padrão tendem a "alucinar" IDs de itens inválidos ou fora de contexto se não forem restringidos.
Ineficiência em Aceleradores (TPUs/GPUs): A solução padrão para impor restrições é usar uma árvore de prefixo (Trie) para mascarar tokens inválidos durante a decodificação. Implementações tradicionais de Trie baseadas em ponteiros (pointer-chasing) causam:
- Latência de Memória: Acessos aleatórios à memória que impedem o coalescing (agrupamento) de acessos, subutilizando a largura de banda de memória de alta velocidade (HBM) de TPUs/GPUs.
- Incompatibilidade de Compilação: Estruturas de controle de fluxo dinâmicas e recursivas são incompatíveis com compiladores de ML estáticos (como o XLA da Google), impedindo a otimização de grafos de computação.

Soluções existentes, como busca binária paralela (PPV), ainda sofrem com escalabilidade logarítmica ( $O(\log |C|)$ ) em relação ao tamanho do conjunto de restrições, tornando-se gargalos em vocabulários massivos (milhões de itens).

2. Metodologia: Framework STATIC

Os autores propõem o STATIC (Sparse Transition Matrix-Accelerated Trie Index for Constrained Decoding), uma técnica que transforma a travessia de grafos (Trie) em operações de matrizes esparsas vetorizadas.

Conceitos Principais:

Flattening do Trie (CSR): A árvore de prefixo é "achatada" e convertida offline em uma Matriz de Transição Esparsa no formato Compressed Sparse Row (CSR).
- Em vez de seguir ponteiros dinâmicos, o sistema realiza uma busca estática em matrizes.
- Isso converte acessos de memória aleatórios em acessos contíguos e vetorizados, ideais para hardware acelerador.
Kernel de Transição de Nós Vetorizado (VNTK): Um algoritmo de decodificação sem ramificações (branch-free) que:
1. Realiza slicing especulativo de um tamanho fixo (baseado no fator de ramificação máximo do nível da árvore).
2. Usa máscaras booleanas para sanitizar os resultados (descartar filhos inexistentes).
3. Mantém o grafo de computação estático, permitindo compilação XLA/Inductor e pipelining completo.
Híbrido Densa-Esparsa: Para os primeiros níveis da árvore (onde a ramificação é alta e densa), utiliza-se uma máscara densa pré-computada para acesso $O(1)$ . Para níveis mais profundos, usa-se a matriz esparsa CSR.

Algoritmo de Decodificação:

O processo de decodificação é reformulado para:

Calcular logits do modelo.
Aplicar a máscara de validade derivada da matriz CSR (ou densa) para definir logits inválidos como $-\infty$ .
Executar a busca em feixe (beam search) apenas sobre os tokens válidos.
Atualizar o estado do nó no Trie usando uma operação de gather vetorizada.

3. Principais Contribuições

Transformação Trie-Matriz: Proposta de mapear restrições de Trie em matrizes esparsas estáticas (CSR), permitindo acesso à memória $O(1)$ e eliminando a sobrecarga de ponteiros.
Algoritmo Nativo para Aceleradores: Design de um algoritmo de decodificação sem ramificações que elimina viagens de ida e volta (round-trips) entre CPU e dispositivo, permitindo execução totalmente no dispositivo (TPU/GPU).
Implantação em Escala Industrial: Sucesso na implantação do sistema no YouTube, restringindo um modelo de recuperação generativa a um vocabulário de 20 milhões de itens frescos.
Melhoria no Cold-Start: Demonstração de que a decodificação constrained pode melhorar significativamente a performance de recomendação para itens novos (cold-start) em benchmarks acadêmicos (Amazon Reviews).

4. Resultados e Desempenho

Eficiência de Sistema (YouTube):

Overhead de Latência: O STATIC adiciona apenas 0,033 ms por passo de decodificação, representando apenas 0,25% do tempo total de inferência.
Aceleração:
- 948x mais rápido que uma implementação de Trie em CPU.
- 47x a 1033x mais rápido que a melhor solução baseada em aceleradores (PPV - Parallel Prefix-Verification).
Escalabilidade: A complexidade de I/O é $O(1)$ em relação ao tamanho do conjunto de restrições, enquanto métodos de busca binária são $O(\log |C|)$ . A latência permanece extremamente baixa mesmo com vocabulários de até 32k tokens e restrições de milhões de itens.
Uso de Memória: Para 20 milhões de itens restritos, o uso de memória HBM é de aproximadamente 1,5 GB, o que é viável para replicação em múltiplos chips.

Impacto no Produto (A/B Testing):

Em um cenário de Home Feed com restrição de "últimos 7 dias":

Aumento de +5,1% nas visualizações de vídeos novos (7 dias).
Aumento de +2,9% nas visualizações de vídeos muito novos (3 dias).
Melhoria na Taxa de Clique (CTR) e satisfação do usuário.

**Desempenho em Cold-Start (Amazon Reviews):**

Em conjuntos de dados de Amazon, o uso de decodificação constrained com o conjunto de itens cold-start melhorou drasticamente o Recall@1 (de 0% para 4,29% no conjunto "Beauty" com 2% de itens novos), superando tanto a busca não restringida quanto o chute aleatório.

5. Significância

Este trabalho é pioneiro ao permitir a primeira implantação em escala de produção de recuperação generativa com restrições estritas.

Quebra de Paradigma: Demonstra que é possível alinhar a flexibilidade dos LLMs com a rigidez necessária para regras de negócios industriais sem sacrificar a latência.
Viabilidade Técnica: Resolve o problema fundamental de como executar estruturas de dados complexas (Tries) em hardware moderno (TPUs/GPUs) que exige grafos estáticos e acesso contíguo à memória.
Impacto Comercial: Permite que grandes plataformas (como o YouTube) utilizem um único modelo generativo grande para múltiplos casos de uso, aplicando filtros de negócio dinâmicos (freshness, estoque, região) diretamente no processo de geração, eliminando a necessidade de filtragem pós-retrieval ineficiente.

Em resumo, o STATIC preenche a lacuna entre a teoria de decodificação constrained e a prática de alto desempenho em hardware, tornando a recuperação generativa uma solução viável e escalável para o futuro dos sistemas de recomendação.