행동할 때와 거절할 때를 배우기: 안전한 다단계 도구 사용을 위한 에이전트 추론 모델 보호

에이전트 AI는 단순 대화형 모델과 달리 실제 도구를 실행하므로 한 번의 실수가 돌이킬 수 없는 피해를 줄 수 있다. 이 논문은 에이전트가 스스로 안전성을 점검하고 위험할 때 명시적으로 거절하는 법을 학습시켜, 복잡한 작업에서도 안전하게 작동하도록 돕는다.

기존 LLM 정렬은 주로 단일 턴의 텍스트 생성에 최적화되어 있어, 여러 단계에 걸쳐 도구를 사용하는 에이전트 환경에서는 한계가 있다. 특히 Attention 메커니즘을 통해 생성된 중간 추론 과정에서 안전 점검이 생략되거나, 유해한 명령이 포함된 도구 결과(Prompt Injection)를 무비판적으로 수용하는 문제가 발생한다.

MOSAIC은 이를 해결하기 위해 안전 점검을 모델의 내재된 추론에만 맡기지 않고 <safety_thoughts>라는 명시적인 토큰 블록으로 분리한다. 이는 모델이 도구를 실행하기 직전에 현재 계획의 위험성, 비가역성, 권한 변경 등을 스스로 되돌아보게 강제하는 일종의 '사고의 브레이크' 역할을 한다.

또한, 학습 과정에서 단순히 최종 결과의 성공 여부만 따지는 것이 아니라, 두 개의 실행 경로(Trajectory)를 비교하여 어떤 경로가 더 안전한 시점에 거절했는지를 평가한다. 이를 통해 모델은 위험이 감지되는 즉시 실행을 멈추는 법을 배우며, 결과적으로 더 적은 토큰을 쓰면서도 높은 안전성을 확보하게 된다.

MOSAIC은 에이전트의 추론 루프를 plan/think → check → act/refuse 단계로 재구성한다. 모델은 먼저 <think> 태그 내에서 계획을 세운 후, 선택적으로 <safety_thoughts> 태그를 호출하여 제안된 도구 호출의 위험성을 평가한다. 평가 결과에 따라 refusal_tool을 호출하여 실행을 즉시 중단하거나, <tool_call>을 통해 실제 동작을 수행한다.

학습에는 GRPO(Group Relative Policy Optimization) 알고리즘을 사용한다. [입력 프롬프트에 대해 n개의 실행 궤적을 샘플링] → [LLM 판독기가 각 궤적 쌍을 비교하여 선호도 행렬을 생성] → [승리 횟수를 기반으로 그룹 내 상대적 보상을 계산] → [이 보상 값을 최대화하는 방향으로 정책 네트워크의 가중치를 갱신]하는 과정을 거친다.

보상 함수 R(τ)는 정렬 보상(ralign), 형식 보상(rfmt), 길이 페널티(plen)의 합으로 구성된다. [궤적의 안전성과 적절성을 입력으로] → [선호도 기반 판독을 통해 0~3 사이의 점수를 산출] → [최종 보상에 합산]하여 모델이 안전하면서도 간결하고 형식을 잘 지키는 응답을 생성하도록 유도한다.

Qwen2.5-7B 모델에서 유해 작업 점수를 50% 감소시켰으며, 유해 작업에 대한 올바른 거절 비율을 87%까지 끌어올렸다. 특히 프롬프트 주입(Prompt Injection) 공격에 대한 거절율이 20% 이상 증가하는 성과를 거두었다.

추론 효율성 측면에서 Qwen3-4B-Thinking 모델은 불필요한 반복 루프를 제거함으로써 전체 토큰 사용량을 4배 이상 절감했다. 안전 점검에 사용되는 토큰은 전체의 20% 미만으로 유지하면서도 안전성을 대폭 강화했다.

PrivacyLens 벤치마크 결과, 개인정보 유출을 최대 23%까지 줄이는 동시에 작업의 유용성(Helpfulness)은 그대로 유지하거나 오히려 향상시켰다. 이는 MOSAIC이 단순히 보수적인 거절만 늘리는 것이 아니라 상황에 맞는 정교한 판단을 내림을 입증한다.

MOSAIC은 에이전트의 행동 공간(Action Space)에 refusal_tool을 일급 시민(First-class citizen)으로 포함시켜, 거절이 단순한 텍스트 생성이 아닌 명시적인 터미널 액션이 되도록 설계했다. 이는 모델이 유해한 중간 도구 호출을 수행하기 전에 실행을 확정적으로 종료할 수 있게 한다.

학습 과정에서 스칼라 보상 모델(Scalar Reward Model) 대신 LLM 판독기(LLM-as-a-judge)를 통한 쌍별 비교(Pairwise Comparison)를 채택했다. 이는 '일찍 거절하는 것'과 '유해한 행동을 다 한 뒤 마지막에 사과하는 것'처럼 최종 상태는 비슷하지만 과정의 안전성이 다른 궤적들을 효과적으로 구분하기 위함이다.

선택적 안전 호출(Selective Safety Invocation) 메커니즘을 통해 모델은 모든 턴에서 안전 점검을 수행하는 대신, 위험이 감지되는 턴에서만 <safety_thoughts>를 생성하도록 학습된다. 이는 추론 비용(Latency)과 안전성 사이의 최적의 트레이드오프를 제공한다.

GRPO를 활용하여 별도의 가치 네트워크(Critic Network) 없이도 그룹 내 상대적 이득을 계산함으로써 학습 안정성을 높였다. 특히 모델 생성 텍스트에만 그래디언트가 흐르도록 도구 출력 토큰을 마스킹 처리하여 에이전트의 의사결정 능력 학습에 집중했다.