BERT-взгляд на оркестрацию больших языковых моделей: токен- и слое-селективные пробы для эффективной однопроходной классификации

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

January 19, 2026

Авторы: Gonzalo Ariel Meyoyan, Luciano Del Corro

cs.AI

Аннотация

Промышленные системы больших языковых моделей (LLM) часто используют отдельные модели для обеспечения безопасности и других задач, требующих интенсивной классификации, что увеличивает задержку, объем видеопамяти (VRAM) и операционную сложность. Мы предлагаем альтернативный подход: повторно использовать вычисления, уже выполненные обслуживающей LLM. Для этого мы обучаем легковесные "зонды" на ее скрытых состояниях и предсказываем метки в рамках того же прямого прохода, который используется для генерации. Мы формулируем задачу классификации как проблему выбора репрезентации из полного тензора скрытых состояний (токен-слой), вместо того чтобы фиксировать конкретный токен или слой (например, логиты первого токена или пулинг последнего слоя). Для реализации этого подхода мы представляем двухуровневый агрегатор, который (i) агрегирует информацию по токенам внутри каждого слоя и (ii) агрегирует полученные сводки по слоям, формируя единое представление для классификации. Мы реализуем этот шаблон с помощью прямого пулинга, 100-тысячного параметрического механизма взвешенного внимания (scoring-attention gate) и зонда на основе понижающей версии много-голового самовнимания (MHA) с числом обучаемых параметров до 35 миллионов. На наборах данных по безопасности и сентимент-анализу наши зонды превосходят подходы, повторно использующие только логиты (например, MULI), и конкурируют со значительно более крупными специализированными моделями, при этом сохраняя задержку, близкую к исходной при обслуживании, и избегая затрат на VRAM и задержку, связанных с использованием отдельного конвейера моделей-фильтров.

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.