StepGuard: Learning Step-Level Guardrails with Scalable Supervision and Safety-Utility Balancing

LLM-based agents can interact with external environments through tool invocation, but this capability also introduces security risks such as file modification, information leakage, and unauthorized actions. Existing guardrails often evaluate completed trajectories, leaving pre-execution monitoring of step-level actions underexplored. We propose StepGuard, a step-level guard model that can audit completed agent trajectories and check tool actions before they are executed. To train StepGuard, we introduce StepGen, an automatic data engine that generates safe and unsafe trajectories with the same context but different actions at the risky step. To further reduce over-defense and under-defense, we propose Balance-GRPO, which dynamically balances learning between safe and unsafe actions based on their observed accuracy. Experiments show that StepGuard achieves the highest average accuracy among open-weight guard models, with performance comparable to GPT-5.4. When used to guard agents on AgentDojo and AgentDyn, StepGuard reduces mean attack success rate by 77.3% relative to the no-guard setting, while mean utility drops by only 2.8 percentage points.

한 줄 요약

StepGuard는 실행 전 단계에서 도구 행동을 점검하는 스텝 레벨 가드 모델로, AgentDojo와 AgentDyn에서 77.3%의 공격 성공률 감소를 달성했다.

핵심 기여도

StepGen이라는 자동 데이터 생성 엔진을 통해 동일한 컨텍스트에서 안전/위험 행동을 비교하는 트래젝토리 생성.
Balance-GRPO 알고리즘을 도입해 안전/위험 행동의 정확도 격차를 기반으로 학습 균형 조정.
StepGuard는 GPT-5.4 수준의 성능을 보이며, 기존 가드 모델 대비 가장 높은 평균 정확도를 달성.
AgentDojo와 AgentDyn에서 공격 성공률(ASR)을 77.3% 감소시키며, 유틸리티는 2.8% 포인트만 감소.

핵심 아이디어

기존 가드레일은 완료된 트래젝토리만 점검하는 데 초점을 맞추어 실행 전 단계의 행동을 모니터링하지 못하는 한계가 있었다. StepGuard는 실행 전 단계에서 도구 행동을 점검하고, 완료된 트래젝토리도 감사하는 스텝 레벨 가드 모델로, 실행 전 보호와 사후 진단을 동시에 수행한다. StepGen은 동일한 컨텍스트에서 단계별로 안전/위험 행동을 생성하여, 대규모 스텝 레벨 감독 데이터를 자동화한다. Balance-GRPO는 GRPO 알고리즘을 확장하여, 안전/위험 행동의 정확도 격차를 기반으로 학습 가중치를 동적으로 조정함으로써 과방어 및 미방어 문제를 완화한다.

기술적 접근법

**StepGen**: ATBench 기반으로 도구와 위험 유형을 샘플링하여, 동일한 실행 프리픽스를 공유하고 위험 단계에서만 분기하는 안전/위험 트래젝토리를 생성.
**Balance-GRPO**: GRPO를 확장한 알고리즘으로, 안전/위험 행동의 정확도 격차를 기반으로 정규화된 이점을 재가중하여 학습 균형을 조정.
**StepGuard 모델**: 4B 파라미터 크기의 가드 모델로, 실행 전 도구 행동 점검과 트래젝토리 감사 기능을 지원.
**하이퍼파라미터**: 학습 데이터는 StepGen을 통해 생성되며, SFT(Instruction Tuning)와 Balance-GRPO를 통해 최적화.

주요 결과

**정확도**: StepGuard는 기존 오픈웨이트 가드 모델 중 가장 높은 평균 정확도를 달성하며, GPT-5.4 수준의 성능을 보임.
**AgentDojo & AgentDyn**: StepGuard는 공격 성공률(ASR)을 77.3% 감소시키며, 유틸리티는 2.8% 포인트만 감소.
**Balance-GRPO 효과**: 기존 GRPO 대비 안전/위험 정확도 격차를 13.0 → 8.0로 줄이며, 유틸리티는 최대 6.7% 포인트 개선, ASR은 0.3% 포인트 증가에 그침.

의의 및 한계

StepGuard는 실행 전 단계에서 행동을 점검함으로써 기존 가드레일의 사후 점검 중심 한계를 극복하며, 실시간 보호와 사후 진단을 결합한 새로운 접근법을 제시한다. Balance-GRPO는 학습 과정에서 안전/유틸리티 균형을 조정함으로써 과방어 및 미방어 문제를 완화하는 데 기여한다. 그러나 AgentHarm 데이터셋에서는 상대적으로 낮은 성능을 보이며, 더 복잡한 위험 상황에 대한 일반화 능력이 제한적이라는 한계가 있다.

실용적 활용

StepGuard는 파일 수정, 정보 유출, 비인가 행동 등이 발생할 수 있는 LLM 기반 에이전트 시스템에 적용 가능하다. 특히, AgentDojo와 AgentDyn과 같은 에이전트 환경에서 실시간 보호와 사후 진단을 동시에 수행할 수 있어, 보안 감사 및 위험 관리에 유용하다.