Uma Visão BERTológica das Orquestrações de LLMs: Sondas Seletivas por Token e Camada para Classificação Eficiente em Passagem Única
A BERTology View of LLM Orchestrations: Token- and Layer-Selective Probes for Efficient Single-Pass Classification
January 19, 2026
Autores: Gonzalo Ariel Meyoyan, Luciano Del Corro
cs.AI
Resumo
Os sistemas de LLM em produção frequentemente dependem de modelos separados para segurança e outras etapas com alta carga de classificação, o que aumenta a latência, o consumo de VRAM e a complexidade operacional. Nós, em vez disso, reutilizamos a computação já paga pelo LLM de serviço: treinamos sondas leves sobre seus estados ocultos e prevemos rótulos no mesmo *forward pass* usado para geração. Enquadramos a classificação como uma seleção de representação sobre o tensor completo de estados ocultos (camada-token), em vez de nos comprometermos com um token fixo ou uma camada fixa (por exemplo, *logits* do primeiro token ou *pooling* da camada final). Para implementar isso, introduzimos um agregador de dois estágios que (i) resume os tokens dentro de cada camada e (ii) agrega os resumos das camadas para formar uma única representação para classificação. Instanciamos este modelo com *pooling* direto, um gate de atenção de pontuação com 100K parâmetros e uma sonda de autoatenção multi-cabeça (MHA) com *downcast* com até 35M parâmetros treináveis. Em *benchmarks* de segurança e sentimento, nossas sondas superam a reutilização baseada apenas em *logits* (por exemplo, MULI) e são competitivas com baselines específicas de tarefa substancialmente maiores, preservando uma latência próxima à de serviço e evitando os custos de VRAM e latência de um *pipeline* separado de modelo de guarda.
English
Production LLM systems often rely on separate models for safety and other classification-heavy steps, increasing latency, VRAM footprint, and operational complexity. We instead reuse computation already paid for by the serving LLM: we train lightweight probes on its hidden states and predict labels in the same forward pass used for generation. We frame classification as representation selection over the full token-layer hidden-state tensor, rather than committing to a fixed token or fixed layer (e.g., first-token logits or final-layer pooling). To implement this, we introduce a two-stage aggregator that (i) summarizes tokens within each layer and (ii) aggregates across layer summaries to form a single representation for classification. We instantiate this template with direct pooling, a 100K-parameter scoring-attention gate, and a downcast multi-head self-attention (MHA) probe with up to 35M trainable parameters. Across safety and sentiment benchmarks our probes improve over logit-only reuse (e.g., MULI) and are competitive with substantially larger task-specific baselines, while preserving near-serving latency and avoiding the VRAM and latency costs of a separate guard-model pipeline.