Eine BERTologie-Perspektive auf LLM-Orchestrierungen: Token- und schichtenselektive Sonden für effiziente Einzeldurchlauf-Klassifikation
A BERTology View of LLM Orchestrations: Token- and Layer-Selective Probes for Efficient Single-Pass Classification
January 19, 2026
papers.authors: Gonzalo Ariel Meyoyan, Luciano Del Corro
cs.AI
papers.abstract
In Produktions-LLM-Systemen werden häufig separate Modelle für Sicherheits- und andere klassifikationsintensive Schritte eingesetzt, was Latenz, VRAM-Bedarf und Betriebskomplexität erhöht. Stattdessen nutzen wir bereits für das servierende LLM durchgeführte Berechnungen erneut: Wir trainieren leichte Proben („Probes“) auf dessen verborgenen Zuständen und sagen Labels im selben Vorwärtsdurchlauf voraus, der auch für die Generierung verwendet wird. Wir formulieren Klassifikation als Repräsentationsauswahl über den gesamten Token-Ebenen-Zustandstensor, anstatt uns auf einen festen Token oder eine feste Ebene (z.B. First-Token-Logits oder Final-Layer-Pooling) festzulegen. Zur Implementierung führen wir einen zweistufigen Aggregator ein, der (i) Token innerhalb jeder Ebene zusammenfasst und (ii) die Ebenen-Zusammenfassungen zu einer einzigen Repräsentation für die Klassifikation aggregiert. Wir instanziieren diese Vorlage mit direktem Pooling, einem Bewertungs-Attentions-Gate mit 100.000 Parametern und einer Probe mit reduziertem multi-head self-attention (MHA) mit bis zu 35 Millionen trainierbaren Parametern. In Sicherheits- und Sentiment-Benchmarks übertreffen unsere Proben die Wiederverwendung nur von Logits (z.B. MULI) und sind mit deutlich größeren aufgabenspezifischen Baseline-Modellen wettbewerbsfähig, bei nahezu unveränderter Servicelatenz und ohne die VRAM- und Latenzkosten einer separaten Guard-Model-Pipeline.
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.