Una Visión BERTológica de las Orquestaciones de LLM: Sondas Selectivas de Tokens y Capas para una Clasificación Eficiente en una Única Pasada

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

Resumen

Los sistemas de LLM en producción a menudo dependen de modelos separados para la seguridad y otros pasos con alta carga de clasificación, lo que incrementa la latencia, la huella de VRAM y la complejidad operativa. Nosotros, en cambio, reutilizamos el cómputo que ya ha sido pagado por el LLM de servicio: entrenamos sondas ligeras sobre sus estados ocultos y predecimos las etiquetas en la misma pasada hacia adelante utilizada para la generación. Enmarcamos la clasificación como una selección de representaciones sobre el tensor completo de estados ocultos (token-capa), en lugar de comprometernos con un token fijo o una capa fija (por ejemplo, los logits del primer token o el pooling de la última capa). Para implementar esto, introducimos un agregador de dos etapas que (i) resume los tokens dentro de cada capa y (ii) agrega los resúmenes de las capas para formar una única representación para la clasificación. Instanciamos esta plantilla con pooling directo, una compuerta de atención puntuadora de 100K parámetros, y una sonda de autoatención multi-cabezal (MHA) de precisión reducida con hasta 35 millones de parámetros entrenables. En diversos benchmarks de seguridad y sentimiento, nuestras sondas mejoran el rendimiento respecto a la reutilización basada únicamente en logits (por ejemplo, MULI) y son competitivas con líneas base específicas de tarea sustancialmente más grandes, mientras preservan una latencia cercana a la del servicio y evitan los costes de VRAM y latencia de un pipeline separado con modelo de guardia.

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.