The Model Knows Which Tokens Matter: Automatic Token Selection via Noise Gating

O artigo apresenta o AutoSelect, um método que reformula a poda de tokens visuais como um problema de comunicação com restrição de capacidade, utilizando um mecanismo de "gate" de ruído e um denoiser para treinar um seletor leve em modelos VLM congelados, permitindo a seleção automática dos tokens mais importantes durante a inferência com ganhos significativos de velocidade e precisão quase total.

O essencial

Imagine que você tem um chef de cozinha genial (o Modelo de Inteligência Artificial) que precisa preparar um prato incrível (responder a uma pergunta sobre uma imagem).

Até agora, o problema era que, antes de cozinhar, o chef recebia uma cesta gigante cheia de ingredientes. A maioria desses ingredientes era apenas casca de banana, papel de embrulho ou pedras (informação redundante). Mas, como o chef não sabia o que era lixo e o que era ouro, ele tinha que examinar cada um dos 576 itens da cesta antes de começar. Isso deixava a cozinha lenta, cansada e cara para operar.

Métodos antigos tentavam resolver isso jogando fora os "lixos" antes de entregar a cesta ao chef. Mas eles faziam isso de forma burra: "Ah, essa pedra parece feia, vou jogar fora". Às vezes, eles jogavam fora algo importante sem querer, ou gastavam tempo demais decidindo o que jogar fora.

O novo método, chamado AutoSelect, propõe uma solução inteligente e elegante. Vamos ver como funciona com uma analogia:

1. O Problema: A "Largura de Banda" Limitada

Pense na comunicação entre a câmera (que tira a foto) e o chef (a IA) como um tubo de água.

• A foto tem muitos detalhes (muitos "tokens" ou pedacinhos de imagem).
• O tubo tem um tamanho limitado. Se você tentar jogar toda a água de um caminhão-pipa num cano de jardim, a pressão explode e o sistema trava.
• O objetivo é: Como enviar apenas a água mais importante pelo tubo, sem entupir?

2. A Solução: O "Porteiro" e o "Filtro de Ruído"

O AutoSelect coloca dois novos funcionários na cozinha, entre a câmera e o chef:

• O Avaliador (Scorer): Ele olha para cada pedacinho da imagem e dá uma nota de "importância".
  • Exemplo: "O olho do gato? Nota 10! O fundo de madeira? Nota 1."
• O Porteiro de Ruído (Noise Gate): Aqui está a mágica. Em vez de simplesmente jogar os itens de nota baixa no lixo (o que quebraria o sistema de aprendizado), o AutoSelect mistura ruído neles.
  • Imagine que os itens importantes são como vinho puro.
  • Os itens sem importância são como água suja.
  • O sistema pega a "água suja" e a mistura com mais água suja (ruído), tornando-a impossível de usar. Já o "vinho puro" fica intacto.
  • Por que fazer isso? Durante o treinamento, o sistema precisa aprender a distinguir o que é importante. Se você apenas joga o lixo fora, o sistema não aprende por que aquilo era lixo. Ao "sujar" o lixo, você força o Avaliador a gritar: "Ei, isso aqui é inútil, não deixe passar!" para que o chef consiga entender a mensagem. É como um jogo de "telefone sem fio" onde só a mensagem clara chega ao final.

3. O "Dedetizador" (Denoiser)

Durante o treinamento, como a "água suja" (ruído) está misturada, o sistema precisa de um ajudante chamado Denoiser. Ele é como um filtro de café que tenta limpar a água suja de volta para o estado original, para que o chef não fique confuso.

• O truque: Esse filtro é "cego". Ele olha para cada gota de água individualmente e não deixa as gotas "conversarem" entre si. Isso impede que a água suja de um lado "contamine" a água limpa do outro.

4. O Resultado Final: A Cozinha Rápida

Depois que o sistema aprendeu (no treinamento), ele não precisa mais do "filtro de detetizador" nem da "mistura de ruído".

• Na hora real (inferência), o Avaliador olha para a imagem, aponta para os 64 melhores ingredientes (dos 576 originais) e diz: "Só leve esses!".
• O resto é descartado instantaneamente.
• O chef recebe apenas os ingredientes essenciais, cozinha 3 vezes mais rápido e o prato fica quase tão bom quanto o original (96,5% da precisão).

Por que isso é revolucionário?

1. Não é "corte e cola": Métodos antigos cortavam pedaços da imagem baseados em regras fixas (como "corte o canto"). O AutoSelect aprende o que é importante para a tarefa específica, como um chef que sabe exatamente quais ervas usar.
2. Rápido e Leve: O processo de escolha leva menos de 1 milissegundo. É como ter um assistente que aponta o dedo para o que importa em um piscar de olhos.
3. Funciona em qualquer lugar: Funciona em modelos pequenos e gigantes, em fotos de baixa ou alta resolução.

Em resumo:
O AutoSelect não joga a informação fora de forma bruta. Ele ensina o modelo a saber o que vale a pena ouvir, transformando um tubo de dados gigante em um canal de alta velocidade, onde apenas o "ouro" passa, e o "lixo" é silenciado. O resultado é uma IA que vê melhor, pensa mais rápido e gasta menos energia.

Resumo técnico

1. O Problema

Os Modelos de Linguagem e Visão (VLMs) modernos, como BLIP-2, InstructBLIP e LLaVA, combinam um codificador visual pré-treinado com um Grande Modelo de Linguagem (LLM). No pipeline padrão, as características visuais (patches de imagem) são projetadas no espaço de embeddings do LLM e tratadas como tokens visuais para decodificação autoregressiva.

• Gargalo Computacional: À medida que as aplicações exigem imagens de alta resolução, múltiplas imagens ou vídeo, o número de tokens visuais aumenta drasticamente. Devido à complexidade quadrática da atenção self-attention em relação ao comprimento da sequência, esses tokens dominam o custo de inferência e a memória (especialmente o cache KV).
• Redundância: Estudos empíricos mostram que a distribuição de atenção é altamente concentrada; muitos tokens recebem atenção quase nula, mas ainda consomem recursos computacionais completos em todas as camadas subsequentes.
• Limitações das Métodos Atuais: As técnicas de poda (pruning) existentes geralmente dependem de sinais proxy locais (como magnitude da atenção ou pontuação de similaridade) e operam como filtros "duros" que descartam tokens irreversivelmente. Isso ignora a questão fundamental de como alocar globalmente a capacidade de representação dentro de um orçamento computacional fixo para maximizar o desempenho.

2. Metodologia: AutoSelect

O artigo propõe o AutoSelect, uma abordagem que reformula a poda de tokens visuais como um problema de aprendizado de representação com restrição de capacidade. Em vez de descartar tokens, o modelo trata a interface entre o codificador visual e o LLM como um canal de largura de banda limitada.

A arquitetura introduz dois módulos leves (Scorer e Denoiser) que são treinados enquanto o VLM base permanece totalmente congelado.

A. Formulação de Canal de Largura de Banda Limitada

O objetivo não é identificar tokens "dispensáveis" para removê-los, mas sim alocar uma capacidade de informação total fixa (orçamento $K$) para os conteúdos mais informativos.

• Treinamento: Nenhum token é removido. Em vez disso, o fluxo de informação de cada token é modulado continuamente.
• Inferência: Apenas os $K$ tokens com as maiores pontuações são selecionados rigidamente (Hard Top-K), eliminando a sobrecarga dos módulos de treinamento.

B. Componentes Principais

1. Scorer (Classificador):

   • Uma rede leve (blocos de Transformer + projeção linear) que atribui uma pontuação de importância escalar a cada token visual.
   • Utiliza um operador Soft Top-K diferenciável para polarizar as pontuações (tornando-as próximas de 0 ou 1) durante o treinamento, permitindo o fluxo de gradientes.
   • A temperatura ($\tau$) é reduzida (annealing) durante o treinamento para que as pontuações se tornem binárias na inferência.

2. Noise Gate (Portão de Ruído) - Preservação de Variância (VP):

   • Para evitar a descontinuidade não diferenciável da remoção direta de tokens, o método injeta ruído nos tokens.
   • A representação do token $i$ é modificada como: $\tilde{x}_i = \sqrt{\alpha_i} x_i + \sqrt{1 - \alpha_i} \epsilon_i$, onde $\alpha_i$ é a pontuação polarizada e $\epsilon_i$ é ruído gaussiano.
   • Mecanismo: Tokens com baixa importância ($\alpha_i \to 0$) são substituídos por ruído isotrópico, enquanto tokens importantes ($\alpha_i \to 1$) são preservados. Isso cria uma restrição de capacidade diferenciável: o modelo só pode passar informação útil se atribuir alta pontuação aos tokens corretos.

3. Denoiser (Dessilenciador) com Atenção Diagonal:

   • Um módulo leve que mapeia a sequência perturbada (com ruído) de volta para a distribuição esperada pelo LLM congelado.
   • Atenção Diagonal: Para evitar que tokens de alta importância "vazem" informação para tokens de baixa importância (contornando a restrição de capacidade), a atenção é restrita a ser diagonal (cada token atende apenas a si mesmo). Isso garante que a restrição de capacidade seja estritamente mantida por token.

C. Fase de Inferência

Durante a inferência, o Noise Gate e o Denoiser são descartados. O Scorer gera as pontuações, e um operador Hard Top-K seleciona os $K$ melhores tokens, mantendo seus índices de posição originais (crucial para embeddings de posição rotativa - RoPE). Isso adiciona latência negligenciável.

3. Principais Contribuições

1. Reformulação Teórica: A poda de tokens é redefinida como aprendizado de representação com restrição de capacidade, modelando a interface como um canal limitado, otimizado apenas com a perda padrão de previsão do próximo token (sem objetivos auxiliares ou anotações extras).
2. Mecanismo de Portão de Ruído (VP Noise Gate): Substitui a decisão binária de "manter ou descartar" por uma modulação contínua de capacidade de informação, permitindo treinamento end-to-end com fluxo total de gradientes e convergindo para seleção Hard Top-K.
3. Eficiência e Generalização: O método alcança alto desempenho com sobrecarga mínima e transfere-se para diferentes arquiteturas de VLM (LLaVA, LLaVA-NeXT, Qwen2.5-VL) sem ajuste específico da arquitetura.

4. Resultados Experimentais

Os testes foram realizados em 10 benchmarks padrão (GQA, MMBench, MME, etc.) e em três arquiteturas principais.

• LLaVA-1.5-7B:
  • Com uma redução de 88,9% dos tokens (mantendo apenas 64 de 576), o AutoSelect retém 96,5% da precisão do modelo completo.
  • Supera métodos de ponta como PRUNESID, HoloV e DART, especialmente em orçamentos de tokens muito restritos.
• LLaVA-NeXT-7B (Alta Resolução):
  • Em imagens de 672x672 (2.880 tokens), mantendo 320 tokens (redução de 88,9%), alcança 96,1% de retenção de desempenho, superando o melhor baseline (HoloV) em 0,4%.
• Qwen2.5-VL-7B:
  • Demonstra generalização para arquiteturas distintas sem tokens [CLS] e com comprimentos de sequência variáveis, superando todos os baselines em todas as taxas de poda.
• Análise de Eficiência:
  • Latência: O módulo de poda do AutoSelect adiciona apenas 0,69 ms de sobrecarga (comparado a 43,39 ms do PRUNESID e 2,77 ms do HoloV).
  • Aceleração: Acelera a fase de "prefill" do LLM em 2,85x, reduzindo o tempo total até o primeiro token (TTFT) de ~150ms para ~72ms (com 64 tokens).
• Validação LLM-Free: Testes de classificação no ImageNet-1K sem o LLM confirmam que a melhoria vem da seleção superior de tokens, e não de adaptação do LLM.

5. Significado e Conclusão

O AutoSelect demonstra que a alocação de capacidade aprendida pode substituir critérios de poda heurísticos. Ao transformar a poda discreta em um problema de otimização contínua via injeção de ruído preservador de variância, o modelo aprende a identificar quais tokens visuais carregam informações relevantes para a tarefa.

A principal vantagem prática é a eficiência extrema: o método oferece ganhos massivos de velocidade e redução de memória com uma sobrecarga de seleção quase nula e sem necessidade de re-treinamento do modelo base ou objetivos de perda complexos. Isso torna a tecnologia viável para implantação em cenários de alta resolução e vídeo, onde o custo de inferência de VLMs é atualmente proibitivo.