[Paper] 정확성 최적화 잔여 활성화 렌즈 (CORAL): 전이 가능하고 보정 인식 추론 시점 스티어링

Source: arXiv - 2602.06022v1

Overview

The paper presents CORAL (Correctness‑Optimized Residual Activation Lens), a lightweight inference‑time technique that nudges large language models (LLMs) toward more accurate and better‑calibrated answers without any extra training. By probing the hidden activations of a model with a regularized MLP, CORAL extracts distributed “correctness signals” and uses them to steer the model’s final prediction, delivering sizable gains in both accuracy and calibration for multiple‑choice QA tasks.

주요 기여

• 실제 정확도에 초점을 맞춘 추론 시점 스티어링을 사용하여 대리 목표(예: 가능도 또는 보상 모델) 대신에.
• 내부 활성화에서 분산된 정확도 정보를 포착하는 Weight‑decay MLP 프로브는 단일 “매직 뉴런”에 의존하지 않습니다.
• 모델에 구애받지 않으며 전이 가능: 동일한 프로브가 서로 다른 7B 파라미터 LLM 3개를 개선하고, 재학습 없이 4개의 보류된 벤치마크에 일반화됩니다.
• 의미 있는 실험적 향상 – 도메인 내 테스트에서 평균 +10 % 정확도와 –50 % 기대 보정 오류(ECE)를 달성하고, 도메인 외 벤치마크에서는 +14 % 정확도와 –49 % ECE를 기록했습니다.
• 연산 효율적인 솔루션: 작은 프로브 네트워크를 몇 번만 순전파하면 되므로 프로덕션 추론 파이프라인에 실용적입니다.

방법론

1. 활성화 스냅샷 수집 – 각 입력(다지선다형 질문)마다 기본 LLM의 여러 층에서 숨겨진 상태를 기록합니다.
2. 정규화된 프로브 학습 – 강한 가중치 감소(L2 정규화)를 적용한 얕은 MLP를 소규모 라벨링된 데이터셋에 학습시켜, 수집된 활성화를 특징으로 사용해 주어진 답 선택이 정답인지 예측합니다. 강한 정규화는 프로브가 개별 뉴런을 암기하기보다 분산된 패턴에 의존하도록 만듭니다.
3. 추론 시 잔차 스티어링 – 모델이 새로운 질문을 처리할 때, 프로브가 각 답 후보의 활성화 스냅샷을 평가하고 “정답 점수”를 생성합니다. 이 점수는 소프트맥스 이전에 모델의 원래 로짓에 잔차로 추가되어, 프로브가 더 옳다고 판단하는 선택지로 순위를 재조정합니다.
4. 보정 인식 조정 – 프로브의 출력이 교차 엔트로피와 같은 적절한 손실로 학습되어 보정되어 있기 때문에, 결과 로짓은 더 나은 신뢰도 추정치를 물려받아 ECE를 감소시킵니다.

전체 파이프라인은 기본 LLM에 대한 그래디언트 업데이트가 전혀 필요 없으며, 작은 프로브에 대한 한 번의 순전파만 수행하면 됩니다.

Results & Findings

• Consistency across architectures – 동일한 프로브 설계가 세 가지 서로 다른 7B‑파라미터 모델(LLaMA‑7B, Falcon‑7B, OpenAI‑스타일 모델 등)에서 모두 작동했습니다.
• Transferability – 하나의 벤치마크(예: ARC‑Easy)에서 훈련된 프로브가 완전히 다른 작업(Math‑MC, HellaSwag)에서도 개선 효과를 보여주었습니다.
• Calibration – Expected Calibration Error가 약 절반으로 감소했으며, 이는 모델의 신뢰도 점수가 실제 정답률과 훨씬 더 잘 일치한다는 의미입니다.

저자들은 이러한 결과를 정답 정보가 다수의 은닉 유닛에 분산되어 있으며, 정규화된 프로브가 이를 신뢰성 있게 추출할 수 있다는 증거로 해석합니다.

Practical Implications

• Plug‑and‑play improvement: 기존 LLM 추론 서비스를 감싸는 얇은 래퍼로 CORAL을 배포; 파인‑튜닝이나 모델 가중치 변경이 필요 없음.
• Cost‑effective scaling: 프로브가 매우 작고(수백 KB) 추론이 약간의 지연만 추가되므로, 대규모 API가 추가 GPU 시간 없이 성능을 향상시킬 수 있음.
• Better user experience: 낮은 ECE는 보다 신뢰할 수 있는 신뢰도 점수로 이어지며, 이는 모델 확률에 기반해 동작하는 하위 시스템(예: 자동 튜터링, 의사결정 지원)에 필수적임.
• Cross‑task robustness: 팀은 적당한 내부 QA 데이터셋으로 단일 프로브를 학습하고, 다양한 하위 MCQA 벤치마크 전반에 걸쳐 이점을 얻을 수 있어 작업별 데이터 수집 필요성을 줄임.
• Safety & alignment: 향상된 캘리브레이션은 과도하게 자신감 있는 환각을 완화하는 데 도움이 되며, 이는 instruction‑tuned LLM에서 흔히 발생하는 실패 모드임.

제한 사항 및 향후 작업

• 다중 선택형 QA에만 적용 – 현재 실험은 MCQA에 초점을 맞추고 있으며, CORAL을 개방형 생성이나 다른 출력 형식으로 확장하는 것은 아직 해결되지 않은 과제입니다.
• 프로브 학습 데이터 요구 – 규모는 작지만, 라벨이 지정된 보정 세트가 필요합니다. 세트가 너무 작거나 도메인이 일치하지 않을 경우 성능이 저하될 수 있습니다.
• 프로브 과적합 가능성 – 강한 가중치 감쇠를 적용하더라도, 프로브가 데이터셋 고유의 특성을 포착할 수 있으므로 실제로 보지 못한 도메인에 대한 체계적인 평가가 필요합니다.
• 저자들이 제시한 향후 연구 방향:
  1. 여러 레이어를 동시에 활용하는 계층적 프로브 탐색.
  2. 잔차 스티어링 개념을 토큰 수준 생성에 적용.
  3. CORAL을 인간 피드백 기반 강화 학습 파이프라인과 통합하여 정확성과 정렬성을 동시에 향상시키기.

저자

• Miranda Muqing Miao
• Young‑Min Cho
• Lyle Ungar

논문 정보

• arXiv ID: 2602.06022v1
• 분류: cs.LG, cs.AI
• 출판일: 2026년 2월 5일
• PDF: PDF 다운로드