A BERTology View of LLM Orchestrations: Token- and Layer-Selective Probes for Efficient Single-Pass Classification

1. Declaração do Problema

Os sistemas de produção modernos de Modelos de Linguagem de Grande Escala (LLM) geralmente dependem de um padrão de orquestração onde um LLM de serviço principal é cercado por modelos auxiliares de "guarda" para moderação de segurança, conformidade com políticas e análise de sentimento. Essa abordagem sofre de três principais desvantagens:

• Latência: Cada solicitação requer uma invocação adicional de modelo (modelo de guarda $\to$ LLM de serviço), dobrando o tempo de inferência.
• Sobrecarga de Recursos: A implantação de modelos separados aumenta a pegada de VRAM e a complexidade operacional.
• Subutilização: O LLM de serviço já computa representações internas ricas (estados ocultos) que contêm sinais de segurança e classificação, mas estes são descartados ao usar apenas os logits de saída finais ou um modelo separado.

Métodos existentes de "reutilização" tentam extrair sinais do modelo de serviço, mas frequentemente se comprometem com pontos de leitura fixos (por exemplo, apenas os logits do primeiro token ou o estado oculto da camada final), potencialmente perdendo informações discriminativas distribuídas por diferentes camadas e posições de token.

2. Metodologia

Os autores propõem um framework de Probe Seletivo por Token e Camada que trata a classificação como um problema de seleção de representações sobre o tensor completo de estados ocultos ($L \times T \times d$, onde $L$=camadas, $T$=tokens, $d$=dimensão). A inovação central é uma Arquitetura de Agregação em Duas Etapas aplicada a um LLM de serviço congelado.

A. Visão Geral da Arquitetura

O probe consiste em um classificador leve treinado sobre os estados ocultos gerados durante a passagem direta (forward pass) do LLM de serviço. Ele não modifica o LLM base.

1. Etapa 1 (Agregação Nível Token): Para cada camada $l$, os $T$ vetores de token são agregados em um único vetor de resumo de camada $v^{(l)}$.
2. Etapa 2 (Agregação Nível Camada): Os $L+1$ vetores de resumo de camada são agregados em um único vetor de representação global $v$.
3. Cabeça de Classificação: Uma camada linear mapeia $v$ para logits de classe.

B. Mecanismos de Agregação

O artigo avalia três mecanismos para realizar a agregação, equilibrando expressividade com eficiência de parâmetros:

1. Pooling Direto: Operadores fixos (Pooling Máximo ou Médio) sobre posições válidas de token/camada. (Menor custo, ~0,003M parâmetros).
2. Portão de Atenção de Pontuação: Um mecanismo aprendível que atribui scores de importância escalar a tokens/camadas usando uma única projeção linear e normalização softmax. (~0,1M parâmetros).
3. Auto-Atenção Multi-Cabeça (MHA): Um módulo MHA reduzido (usando redução agressiva de dimensão, por exemplo, $d/16$) para capturar interações complexas entre tokens e camadas. (~35M parâmetros).

C. Fluxo de Trabalho Operacional

Durante a inferência, o probe executa em paralelo com o processo de geração. Se o probe prever um rótulo "inseguro", a camada de orquestração pode interromper imediatamente a geração antes que qualquer token seja transmitido, retornando uma recusa ou re-promptando o modelo. Isso elimina a necessidade de uma segunda chamada de modelo.

3. Contribuições Principais

• Mudança de Leituras Fixas para Seletivas: O artigo avança além de leituras de camada fixa ou token fixo (como [CLS] do BERT ou logits do primeiro token) para um mecanismo de seleção aprendida que identifica onde no tensor $L \times T$ o sinal é mais separável.
• Modelo de Agregação em Duas Etapas: Uma arquitetura inovadora que primeiro resume tokens dentro das camadas e depois agrega através das camadas, permitindo que o modelo aprenda a profundidade específica da tarefa e a importância dos tokens.
• Compromisso entre Eficiência e Desempenho: Demonstra que classificação de alta precisão pode ser alcançada com parâmetros treináveis mínimos (até 35M) e latência de inferência adicional zero em comparação com o modelo de serviço isolado.
• Generalização: Valida a abordagem em diversas arquiteturas, incluindo modelos densos (Llama-3.2-3B) e Mistura de Especialistas (GPT-OSS-20B, Qwen3-30B-A3B).

4. Resultados Experimentais

Moderação de Segurança (ToxicChat & WildGuardMix)

• Desempenho: O probe MHA alcançou 84,51 F1 no ToxicChat e 88,55 F1 no WildGuardMix.
• Comparação:
  • Superou as bases de reutilização apenas de logits (por exemplo, MULI: 77,8 F1).
  • Igualou ou superou modelos de guarda independentes (por exemplo, ToxicChat-T5-large: 82,2 F1; WildGuard: 88,9 F1) apesar de usar ordens de magnitude menos parâmetros (35M vs. 780M–8B).
  • Latência Zero Extra: Diferentemente de modelos de guarda que exigem uma chamada de API separada, o probe executa na mesma passagem direta.
• Cross-Arquitetura: A classificação de desempenho (Pooling < Atenção de Pontuação < MHA) permaneceu consistente entre os modelos Llama, GPT-OSS e Qwen.

Sentimento e Emoção (IMDB, SST-2, Emotion)

• Os probes superaram significativamente o prompting zero-shot e few-shot (por exemplo, 87,68% de precisão em Emotion vs. 64,05% para MULI).
• O desempenho foi competitivo com grandes classificadores específicos de tarefa (por exemplo, DeBERTa, RoBERTa) enquanto mantinha um pipeline de inferência de passada única.

Análise de Eficiência

• Latência: O probe MHA adicionou apenas ~14ms de latência por amostra em comparação com o modelo base, enquanto um pipeline guarda-então-serve adicionou ~60–100ms.
• Memória: As variantes do probe mantiveram o pico de memória GPU próximo à pegada do modelo base (~6,5–7,0 GB), enquanto pipelines de guarda exigiram 8,0–22,8 GB.

Interpretabilidade (Atenção de Camada)

• A análise do portão de Atenção de Pontuação revelou que sinais de segurança estão distribuídos. Prompts tóxicos tendiam a ativar camadas posteriores (L17–L28), enquanto prompts não tóxicos dependiam das camadas finais e embeddings. Classificações incorretas frequentemente mostraram padrões de atenção alinhados com a classe prevista em vez da verdade fundamental, sugerindo que o probe aprende variedades de características distintas para diferentes tipos de conteúdo.

5. Significado e Conclusão

Este artigo fornece uma alternativa convincente ao padrão de orquestração "modelo de guarda" padrão. Ao aproveitar os princípios da BERTology (que diferentes camadas codificam diferentes abstrações) e aplicá-los a LLMs modernos apenas decodificadores, os autores mostram que:

1. Tarefas de segurança e NLU já estão "resolvidas" internamente pelo modelo de serviço; o desafio é extrair o sinal de forma eficiente.
2. Classificação em passada única é viável para produção, oferecendo uma redução massiva em latência e custos de infraestrutura sem sacrificar a precisão.
3. Agregação aprendida (especificamente MHA) é superior ao pooling fixo, provando que a "melhor" camada para uma tarefa não é estática, mas depende do conteúdo de entrada.

O trabalho sugere um futuro onde a orquestração de LLMs é simplificada, reduzindo a necessidade de pipelines de segurança pesados e separados e permitindo sistemas de IA mais responsivos e economicamente eficientes.