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

1. Declaración del Problema

Los sistemas de producción modernos de Modelos de Lenguaje Grandes (LLM) suelen depender de un patrón de orquestación donde un LLM de servicio principal está rodeado por modelos auxiliares de "guardia" para moderación de seguridad, cumplimiento de políticas y análisis de sentimientos. Este enfoque adolece de tres desventajas principales:

• Latencia: Cada solicitud requiere una invocación adicional de modelo (modelo de guardia $\to$ LLM de servicio), duplicando el tiempo de inferencia.
• Sobrecarga de Recursos: Desplegar modelos separados aumenta la huella de VRAM y la complejidad operativa.
• Subutilización: El LLM de servicio ya calcula representaciones internas ricas (estados ocultos) que contienen señales de seguridad y clasificación, pero estas se descartan al utilizar solo los logits de salida final o un modelo separado.

Los métodos de "reutilización" existentes intentan extraer señales del modelo de servicio, pero a menudo se comprometen a puntos de lectura fijos (por ejemplo, solo los logits del primer token o el estado oculto de la capa final), perdiendo potencialmente información discriminativa distribuida a través de diferentes capas y posiciones de tokens.

2. Metodología

Los autores proponen un marco de Sonda Selectiva por Token y Capa que trata la clasificación como un problema de selección de representaciones sobre el tensor completo de estados ocultos ($L \times T \times d$, donde $L$=capas, $T$=tokens, $d$=dimensión). La innovación central es una Arquitectura de Agregación en Dos Etapas aplicada a un LLM de servicio congelado.

A. Descripción General de la Arquitectura

La sonda consiste en un clasificador ligero entrenado sobre los estados ocultos generados durante el paso hacia adelante del LLM de servicio. No modifica el LLM base.

1. Etapa 1 (Agregación a Nivel de Token): Para cada capa $l$, los $T$ vectores de tokens se agregan en un único vector de resumen de capa $v^{(l)}$.
2. Etapa 2 (Agregación a Nivel de Capa): Los $L+1$ vectores de resumen de capa se agregan en un único vector de representación global $v$.
3. Cabeza de Clasificación: Una capa lineal mapea $v$ a los logits de clase.

B. Mecanismos de Agregación

El artículo evalúa tres mecanismos para realizar la agregación, equilibrando expresividad con eficiencia de parámetros:

1. Agrupación Directa (Pooling): Operadores fijos (Max o Mean pooling) sobre posiciones válidas de tokens/capas. (Menor costo, ~0.003M parámetros).
2. Puerta de Puntuación por Atención: Un mecanismo aprendible que asigna puntuaciones de importancia escalar a tokens/capas utilizando una proyección lineal única y normalización softmax. (~0.1M parámetros).
3. Autoatención Multi-Cabeza (MHA): Un módulo MHA reducido (usando reducción agresiva de dimensión, por ejemplo, $d/16$) para capturar interacciones complejas entre tokens y capas. (~35M parámetros).

C. Flujo de Trabajo Operativo

Durante la inferencia, la sonda se ejecuta en paralelo con el proceso de generación. Si la sonda predice una etiqueta "insegura", la capa de orquestación puede detener inmediatamente la generación antes de que se transmita cualquier token, devolviendo una negativa o re-promptando al modelo. Esto elimina la necesidad de una segunda llamada al modelo.

3. Contribuciones Clave

• Cambio de Lecturas Fijas a Selectivas: El artículo avanza más allá de las lecturas de capa fija o token fijo (como [CLS] de BERT o logits del primer token) hacia un mecanismo de selección aprendido que identifica dónde en el tensor $L \times T$ la señal es más separable.
• Plantilla de Agregación en Dos Etapas: Una arquitectura novedosa que primero resume los tokens dentro de las capas y luego agrega a través de las capas, permitiendo que el modelo aprenda la importancia específica de la tarea en profundidad y token.
• Compensación entre Eficiencia y Rendimiento: Demuestra que se puede lograr una clasificación de alta precisión con parámetros entrenables mínimos (hasta 35M) y cero latencia de inferencia adicional en comparación con el modelo de servicio por sí solo.
• Generalización: Valida el enfoque en diversas arquitecturas, incluidos modelos densos (Llama-3.2-3B) y de Mezcla de Expertos (GPT-OSS-20B, Qwen3-30B-A3B).

4. Resultados Experimentales

Moderación de Seguridad (ToxicChat & WildGuardMix)

• Rendimiento: La sonda MHA logró 84.51 F1 en ToxicChat y 88.55 F1 en WildGuardMix.
• Comparación:
  • Superó las líneas base de reutilización solo de logits (por ejemplo, MULI: 77.8 F1).
  • Igualó o superó a los modelos de guardia independientes (por ejemplo, ToxicChat-T5-large: 82.2 F1; WildGuard: 88.9 F1) a pesar de utilizar órdenes de magnitud menos parámetros (35M vs. 780M–8B).
  • Cero Latencia Extra: A diferencia de los modelos de guardia que requieren una llamada a API separada, la sonda se ejecuta en el mismo paso hacia adelante.
• Cruz-Arquitectura: La clasificación de rendimiento (Agrupación < Puerta de Puntuación por Atención < MHA) se mantuvo consistente en los modelos Llama, GPT-OSS y Qwen.

Sentimiento y Emoción (IMDB, SST-2, Emotion)

• Las sondas superaron significativamente a la indicación (prompting) cero-shot y pocos-shot (por ejemplo, 87.68% de precisión en Emotion frente a 64.05% para MULI).
• El rendimiento fue competitivo con clasificadores grandes específicos de la tarea (por ejemplo, DeBERTa, RoBERTa) mientras se mantenía un pipeline de inferencia de un solo paso.

Análisis de Eficiencia

• Latencia: La sonda MHA añadió solo ~14ms de latencia por muestra en comparación con el modelo base, mientras que un pipeline de guardia-antes-de-servicio añadió ~60–100ms.
• Memoria: Las variantes de la sonda mantuvieron la memoria pico de GPU cerca de la huella del modelo base (~6.5–7.0 GB), mientras que los pipelines de guardia requirieron 8.0–22.8 GB.

Interpretabilidad (Atención de Capa)

• El análisis de la puerta de puntuación por atención reveló que las señales de seguridad están distribuidas. Los prompts tóxicos tendían a activar capas posteriores (L17–L28), mientras que los prompts no tóxicos dependían de las capas finales y los embeddings. Las clasificaciones incorrectas a menudo mostraron patrones de atención alineados con la clase predicha en lugar de la verdad fundamental, lo que sugiere que la sonda aprende variedades de características distintas para diferentes tipos de contenido.

5. Significado y Conclusión

Este artículo proporciona una alternativa convincente al patrón de orquestación estándar de "modelo de guardia". Al aprovechar los principios BERTológicos (que diferentes capas codifican diferentes abstracciones) y aplicarlos a LLM modernos de solo decodificador, los autores muestran que:

1. Las tareas de seguridad y comprensión del lenguaje natural (NLU) ya están "resueltas" internamente por el modelo de servicio; el desafío es extraer la señal de manera eficiente.
2. La clasificación en un solo paso es viable para producción, ofreciendo una reducción masiva en latencia y costos de infraestructura sin sacrificar la precisión.
3. La agregación aprendida (específicamente MHA) es superior a la agrupación fija, demostrando que la "mejor" capa para una tarea no es estática, sino que depende del contenido de entrada.

El trabajo sugiere un futuro donde la orquestación de LLM se simplifica, reduciendo la necesidad de pipelines de seguridad pesados y separados y permitiendo sistemas de IA más receptivos y rentables.