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이 이미 수행한 계산을 재활용합니다: 은닉 상태에 경량 프로브를 학습시켜 생성에 사용된 동일한 순전파 과정에서 레이블을 예측합니다. 분류를 고정된 토큰이나 고정된 계층(예: 첫 번째 토큰 로짓 또는 최종 계층 풀링)에 의존하기보다 전체 토큰-계층 은닉 상태 텐서에 대한 표현 선택 문제로 재구성합니다. 이를 구현하기 위해 (i) 각 계층 내 토큰을 요약하고 (ii) 계층별 요약을 종합하여 분류를 위한 단일 표현을 형성하는 2단계 집계기를 도입합니다. 우리는 이 템플릿을 직접 풀링, 100K 매개변수 스코어링 어텐션 게이트, 최대 35M의 학습 가능 매개변수를 가진 다운캐스트 다중 헤드 자가 어텐션(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.