ProbGuard: Calibrated Safety Risk Estimation from LLM Output Distributions

당신은 마술 쇼를 보고 있다고 상상해 보세요. 그곳의 마술사는 이야기를 쓰거나, 수학 문제를 풀거나, 심지어 친구처럼 대화까지 할 수 있는 아주 똑똑한 로봇입니다. 이 로봇은 '거대 언어 모델(Large Language Model, LLM)'이라고 불립니다. 이 로봇은 믿기지 않을 정도로 재능이 뛰어나지만, 강력한 도구에는 항상 어두운 면이 따르듯, 때로는 폭탄을 만드는 법이나 누군가를 해치는 법과 같이 위험하거나 못되거나 불법적인 내용을 말하도록 속임을 당할 수 있습니다. 이는 현실 세계에서 이 로봇들을 사용하는 모든 이들에게 큰 걱정거리입니다.

로봇이 나쁜 아이디어를 쏟아내는 것을 막기 위해, 우리는 보통 '가드레일(guardrail)'을 설치합니다. 가드레일을 클럽 입구를 지키는 엄격한 보안 요리라고 생각해 보세요. 과거에 이 보안 요원은 로봇이 쓴 완성된 문장만을 확인했습니다. 만약 문장이 나빠 보이면, 보안 요력은 그것을 쫓아냈습니다. 하지만 문제가 있었습니다. 보안 요력이 너무 느렸던 것입니다. 로봇이 문장을 다 마칠 때쯤이면 이미 피해는 발생한 후였습니다. 또한, 보안 요력은 너무 흑백논리적이었습니다. 로봇이 다음에 무슨 말을 할지 '확신하지 못하는' 상태라는 점을 이해하지 못했습니다. 그것은 마치 보안 요력이 최종 티켓만 검사하고, 로봇이 말을 결정하는 동안 손을 떨며 불안해하던 모습은 무시하는 것과 같았습니다.

여기서 ProbGuard라는 새로운 아이디어가 등장합니다. ProbGuard는 단순히 완성된 문장을 보는 대신, 로봇의 생각이 진행되는 동안 그 '사고 과정'을 관찰하는 매우 예리한 탐정처럼 행동합니다. ProbGuard는 로봇이 이미 타이핑한 단어들만 보는 것이 아니라, 다음에 어떤 단어를 말할지에 대한 로봇의 '직감'을 들여다봅니다. 만약 로봇이 몇 초 안에 위험한 말을 할 확률이 90%라고 판단되면, ProbGuard는 문장이 끝나기도 전에 즉시 비상 브레이크를 밟아 로봇을 멈춰 세웁니다. 이것은 자동차가 사고가 난 후에 경찰을 부르는 것이 아니라, 차가 경로를 이탈하는 것을 보고 즉시 멈추는 것과 같습니다.

기존 가드레일의 문제점

이 논문은 기존의 방식이 '단어 맞히기' 게임과 같으며, 오직 최종 정답만을 볼 수 있다고 설명합니다. 대부분의 현재 안전 시스템은 단순한 분류기(classifier)처럼 작동합니다. 즉, 완성된 문장을 가져와서 "안전함" 또는 "위험함"이라고 판정합니다.

저자들은 여기에는 두 가지 큰 결함이 있다고 주장합니다:

1. 너무 늦습니다: 문장이 끝났을 때는 이미 나쁜 결과물을 막기에 너무 늦은 시점입니다.
2. '아마도'를 무시합니다: 특히 로봇이 아직 생각 중일 때, 안전은 항상 예/아니오의 문제가 아닙니다. 로봇은 무해해 보이는 문장으로 시작했지만 쉽게 위험한 문장으로 변할 수도 있습니다. 기존 시스템은 로봇의 내부적인 '불확실성'을 버리고 오직 최종 텍스트만을 봅니다. 로봇이 망설였거나, 혹은 안전하거나 위험한 여러 갈래의 길을 고려하고 있었다는 사실을 무시하는 것입니다.

ProbGuard의 작동 방식: 수정구슬 접근법

ProbGuard는 안전을 단순한 라벨이 아닌 '확률'로 취급함으로써 게임의 판도를 바꿉니다. 로봇이 갈림길에 서 있다고 상상해 보세요. 기존의 가드레일은 로봇이 길을 선택하고 그 길을 따라 걸어간 뒤에야 그 길이 절벽으로 이어지는지 확인합니다. 하지만 ProbGuard는 로봇이 그 자리에 서 있는 동안 로봇이 들고 있는 '지도'를 봅니다.

ProbGuard가 사용하는 단계별 마술 트릭은 다음과 같습니다:

1. "만약에" 시뮬레이션 (몬테카를로 샘플링)
로봇의 현재 생각이 얼마나 위험한지 알기 위해, ProbGuard는 다음과 같은 질문을 던집니다. "만약 이 시점에서 로봇이 계속 말을 하게 둔다면, 나쁜 말을 할 확률은 얼마인가?"
미래를 완벽하게 예측할 수는 없기에, ProbGuard는 머릿속에서 이 시나리오를 수천 번 실행합니다. 이 시나리오를 16가지 다른 방식으로 완성하도록 시뮬레이션합니다(주사위를 16번 굴려 확률을 보는 것과 같습니다). 16번 중 15번은 로봇이 안전한 말을 하지만 1번은 위험한 말을 한다면, ProbGuard는 작은 위험이 있음을 인지합니다. 만약 16번 중 10번이 나쁜 말을 한다면, 위험이 높은 것입니다. 이를 통해 ProbGuard는 단순한 "안전/위험" 라벨 대신 정밀한 "위험 점수"(0과 1 사이의 숫자)를 산출합니다.

2. 다음 단어의 "유령" 읽기
로봇이 텍스트를 생성할 때, 로봇은 단 하나의 단어만 고르는 것이 아니라 다음에 말할 수 있는 모든 가능한 단어에 대한 확률을 계산합니다. 예를 들어, 로봇이 "하늘은..."이라고 말하려 할 때, 로봇은 "푸르다"(30% 확률), "초록색이다"(5% 확률), "불타고 있다"(2% 확률)와 같이 생각할 수 있습니다.
기존 시스템은 보통 확률이 가장 높은 단어인 "푸르다"만을 봅니다. 하지만 ProbGuard는 전체 가능성의 목록을 봅니다. "불타고 있다"가 아주 낮은 확률을 가지고 있더라도, 만약 그 작은 확률이 재앙으로 이어진다면 그것은 중요합니다. ProbGuard는 이 전체 확률 목록을 특수한 코드("확률 가중 표현")로 변 변환하여, 로봇의 비밀스러운 뇌(숨겨진 상태, hidden states)를 들여다보지 않고도 읽을 수 있게 합니다. 이 덕분에 ProbGuard는 특정 브랜드의 로봇에 국한되지 않고 어떤 종류의 로봇과도 함께 작동할 수 있습니다.

3. 조기 중단 (Early Stop)
ProbGuard는 위험 점수를 계산한 후 즉시 행동할 수 있습니다. 논문은 ProbGuard가 로봇이 입력하는 단 첫 10단어(또는 디코딩 단계)만 보고도 그것이 위험해질 것인지 결정할 수 있음을 보여줍니다. 위험 점수가 너무 높으면, ProbGuard는 로봇이 문장을 완성하기 전 바로 멈춰 세웁니다.

숫자가 말해주는 것들

연구진은 ProbGuard를 Llama-Guard3나 ShieldGemma와 같은 현재 최고의 가드들을 포함한 13개의 다른 안전 방법들과 테스트했습니다. 그들은 누가 가장 잘 예측하는지 알아보기 위해 세 가지 유형의 로봇과 세 가지 세트의 위험한 질문을 사용했습니다.

• 교정(Calibration): 이것은 "로봇이 자신의 위험에 대해 얼마나 정직한가"를 뜻하는 멋진 단어입니다. 만약 ProbGuard가 위험 확률이 90%라고 말한다면, 실제로 90%의 경우에 위험이 발생하는가? 연구 결과, ProbGuard는 다른 누구보다 이 부분에서 훨씬 뛰어났습니다. 기존의 가장 좋은 방법과 비교했을 때 위험 예측 오류를 약 79.6% 줄였습니다.
• 공격 차단: 팀은 로봇이 안전 규칙을 무시하도록 만드는 특수 기술인 여섯 가지 "탈옥(jailbreak)" 기법을 사용하여 로봇을 속이려 했습니다. ProbGuard는 놀라울 정도로 효과적이었습니다. 단 10단어만 보고도 공격을 거의 완벽하게 차단하여, 이러한 기술의 성공률을 최대 1% 미만으로 제한했습니다. 이에 비해 기존의 가장 좋은 가드레일은 약 2.4%의 공격을 허용했습니다.
• 속도: ProbGuard는 또한 빠릅니다. 1,000개의 샘플을 약 36.4초 만에 확인할 수 있으며, 이는 다른 무거운 안전 시스템들보다 약 52.7% 더 빠른 속도입니다.

핵심 요약

논문은 결론적으로 ProbGuard가 안전을 단순한 "예/아니오" 체크리스트로 취급하지 않기 때문에 중요한 진전이라고 설명합니다. 대신, 안전을 로봇이 생각하는 동안 측정하고 관리할 수 있는 유동적이고 변화하는 확률로 취급합니다.

로봇의 "직감"(출력 분포)을 사용하고 많은 가능한 미래를 시뮬레이션함으로써, ProbGuard는 위험한 생각이 온전한 문장이 되기 전에 잡아낼 수 있습니다. 이 방식은 다양한 로봇 모델에서 작동하며, 로봇의 비밀스러운 뇌를 들여다볼 필요가 없고, 높은 정확도와 속도를 모두 갖추고 있습니다. 저자들은 이 접근 방식이 우리가 실수를 저지르는 바로 그 첫 순간에 AI 시스템을 멈출 수 있게 함으로써, 재앙이 완전히 펼쳐질 때까지 기다리는 것이 아니라 사고를 미연에 방지할 수 있게 해준다고 제안합니다.