문맥 위험을 제어하는 조기 종료 기반 안전 LLM 프레임워크

초록

본 논문은 사용자가 제공하는 컨텍스트가 해롭거나 부정확할 경우, 대형 언어 모델(LLM)의 성능 저하를 방지하기 위해 “제로샷” 성능을 안전 기준으로 삼고, 동적 조기 종료와 분포 자유 위험 제어(DFRC)를 결합한 방법을 제안한다. 조기 종료 시점에서 신뢰도가 기준값을 넘지 않으면 컨텍스트를 무시하고 제로샷 예측으로 돌아가며, 새로운 컨텍스트‑인식 손실을 정의해 위험을 정량화한다. Learn‑then‑Test 프레임워크를 변형해 음수 손실도 보존하면서 위험을 제어하고, 9개의 다양한 작업과 5개의 모델에 대해 50% 이상 속도 향상과 해로운 컨텍스트에 대한 안전성을 입증한다.

상세 분석

이 연구는 LLM이 사용자 제공 컨텍스트에 과도하게 의존해 “오버싱킹(overthinking)” 현상을 보이는 문제를 근본적으로 해결하고자 한다. 먼저, 컨텍스트가 전혀 없을 때의 제로샷 성능을 “안전 기준(baseline)”으로 정의한다. 이는 사전 배포 단계에서 광범위한 안전 검증을 거친 모델이므로, 위험 제어의 기준점으로 적합하다. 이후, 분포 자유 위험 제어(DFRC) 이론을 적용해, 컨텍스트가 모델 성능을 이 기준 이하로 떨어뜨리는 확률을 사용자가 지정한 허용 수준 ε 이하로 제한한다. 핵심 메커니즘은 동적 조기 종료(dynamic early exit)이다. 각 레이어 l에서 예측 확률의 최대값을 신뢰도 Cₗ로 사용하고, 사전 정의된 임계값 λ와 비교한다. Cₗ ≥ λ이면 해당 레이어에서 즉시 출력을 반환하고, 그렇지 않으면 다음 레이어로 진행한다. 모든 레이어를 통과해도 λ를 초과하지 않으면 컨텍스트를 완전히 무시하고 제로샷 모델(p_L(·|x))의 출력을 사용한다(식 2). 이렇게 하면 해로운 컨텍스트가 깊은 레이어에서 모델을 오염시키는 것을 방지하면서, 유용한 컨텍스트가 존재할 경우 조기 종료을 통해 계산 비용을 절감하고 성능 이득을 얻을 수 있다.

새로운 컨텍스트‑인식 손실 ℓ_c(λ;x,y,c) = ℓ(ȳ_λ(x,c),y) – ℓ(ŷ(x),y) 를 도입한다. 여기서 ȳ_λ는 조기 종료 후 컨텍스트를 활용한 예측, ŷ는 제로샷 예측이다. 해로운 컨텍스트에서는 ℓ_c > 0(위험 증가), 유용한 컨텍스트에서는 ℓ_c < 0(위험 감소)로 나타나며, 이는 기존 조기 종료 연구가 최종 레이어와의 차이만을 고려해 오버싱킹을 포착하지 못했던 점을 보완한다.

ℓ_c는 λ에 대해 비단조(monotonic)하지 않으므로, 일반적인 위험 제어 기법(예: conformal prediction)으로는 적용이 어렵다. 저자들은 Learn‑then‑Test(LTT) 프레임워크를 채택하고, 손실이