LLMオーケストレーションのBERTology的視点：効率的な単一パス分類のためのトークン・層選択的プローブ

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が既に行った計算を再利用する手法を提案します。具体的には、LLMの隠れ状態に軽量なプローブを接続して訓練し、生成と同じフォワードパス内でラベル予測を同時に行います。分類タスクを、特定のトークン（例：先頭トークンのロジット）や特定の層（例：最終層のプーリング）に固定するのではなく、全トークン・全層にわたる隠れ状態テンソルからの表現選択問題として再定義します。これを実現するため、2段階のアグリゲータを導入しました。(i) 各層内のトークン情報を要約し、(ii) 層ごとの要約を統合して、分類のための単一の表現を形成します。この枠組みを具体化するため、直接プーリング、10万パラメータのスコアリングアテンションゲート、最大3500万の学習可能パラメータを持つダウンキャスト型マルチヘッド自己注意（MHA）プローブを実装しました。安全性評価と感情分析のベンチマークにおいて、提案するプローブはロジットのみの再利用手法（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.