개요

에이전트 RL 학습 루프를 멀티턴으로 확장하는 순간 이상한 현상이 나타난다. 손실이 이유 없이 치솟고, 가끔 shape mismatch 오류가 난다. 원인은 거의 항상 같다 — 재토크나이제이션이 모델이 실제 샘플링한 토큰을 슬쩍 바꿔치기한다. HuggingFace의 Quentin Gallouédec과 Kashif Rasul이 이 문제의 원인과 해법을 TITO(Token-In, Token-Out) 원칙으로 정리했다.

핵심 요약

  • MITO 버그: 메시지 리스트를 매 턴 재렌더링·재토크나이즈하면 모델이 실제 생성한 토큰과 다른 시퀀스가 나올 수 있다
  • 원인: BPE 토크나이제이션은 비가역적 — 디코딩 후 재인코딩하면 다른 token ID가 나올 수 있다
  • 결과: 그래디언트가 모델이 실제 생성하지 않은 토큰에 걸림 → 신뢰할 수 없는 학습 신호
  • TITO 원칙: 디코딩한 토큰은 절대 다시 인코딩하지 않는다. 실행 중 토큰 버퍼를 단일 진실 소스로 유지
  • 도구 응답 처리: 템플릿 diff(prefix vs full render) 한 번으로 delta를 계산해 버퍼에 추가

문제: 재토크나이제이션이 토큰을 바꾼다

RL의 핵심 규칙은 단순하다. 그래디언트는 모델이 샘플링한 정확한 토큰에만 걸려야 한다. 단일 턴에서는 자명하다. 멀티턴에서 이 규칙이 깨진다.

자연스럽게 짜는 멀티턴 루프는 이렇게 생겼다.

  1. 메시지 리스트 유지
  2. 매 턴: 메시지 리스트 렌더링 → 토크나이즈 → 모델 생성 → 파싱
  3. 도구 호출이면 실행 후 결과 추가
  4. 롤아웃 끝나면 전체 대화를 다시 토크나이즈 → 역전파

이 루프의 핵심 결함은 4번 단계다. 롤아웃 마지막에 대화 전체를 다시 토크나이즈하면 모델이 실제 생성한 토큰 시퀀스와 미묘하게 달라질 수 있다.

이유는 BPE 토크나이제이션의 비가역성이다. 동일한 문자열이라도 BPE 세그멘테이션이 여러 가지 존재한다. 그 위에 JSON 직렬화 공백, 인자 순서, 불리언 표기(false vs False), 특수 토큰 재렌더링 방식이 추가적인 불확정성을 더한다.

코드는 정상 실행된다. 하지만 수학이 조용히 망가진다. 그래디언트가 모델이 생성하지 않은 토큰에 걸린다.

TITO 원칙: 디코딩한 것은 다시 인코딩하지 않는다

해법은 하나의 규칙이다.

디코딩한 토큰은 절대 다시 인코딩하지 않는다.

구현은 이렇다. 모델이 생성한 토큰들을 메시지 리스트 대신 실행 중 토큰 버퍼에 누적한다. 이 버퍼가 단일 진실 소스다.

파싱은 여전히 한다. 도구 호출인지 확인하려면 디코딩이 필요하다. 하지만 파싱 결과는 라우팅에만 쓴다 — 다음 프롬프트로 절대 피드백하지 않는다. 디코드 → 라우팅 → 버린다.

이 단일 변경으로 두 가지 문제가 사라진다.

Per-turn 경계 손실 문제: 기존 방식은 마지막에 재렌더링된 문자열에서 역방향으로 경계를 추출했다. TITO는 버퍼가 자라면서 경계를 실시간으로 추적한다. 어시스턴트 토큰인지 도구 응답 토큰인지는 추가된 시점에 이미 알고 있다. 손실 마스크를 사후 복원할 필요가 없다.

토큰 드리프트 문제: 버퍼는 재인코딩되지 않는다. 정책이 샘플링한 토큰이 그래디언트 아래 놓이는 토큰과 완전히 일치한다.

도구 응답을 버퍼에 추가하는 방법

TITO 루프에서 템플릿을 사용하는 부분은 딱 하나 남는다 — 도구 응답 토크나이제이션이다.

도구가 결과를 반환하면, 그 결과를 모델이 기대하는 형식(특수 토큰 래핑, 역할 마커 등)으로 인코딩해서 버퍼에 추가해야 한다. 이걸 “템플릿 diff” 방식으로 구현한다.

도구 호출 직전까지의 메시지로 렌더링 한 번, 도구 응답을 포함해 렌더링 한 번. 두 결과의 차이(delta)가 정확히 도구 응답 래핑 토큰들이다. 이 delta만 버퍼에 추가한다.

prefix = tok.apply_chat_template(messages_prefix, return_dict=False)
full   = tok.apply_chat_template(messages_full, return_dict=False, add_generation_prompt=True)
delta  = full[len(prefix):]  # 이것만 버퍼에 추가

이 연산이 유효하려면 채팅 템플릿이 prefix-preserving이어야 한다 — 새 메시지를 추가해도 앞의 토큰이 바뀌지 않는 성질이다. 실제로 대부분의 오픈소스 모델 템플릿이 이 조건을 만족한다.

두 가지 해법 비교

이 문제를 해결하는 접근이 두 가지다.

Renderers 방식 (PrimeIntellect): 모델 패밀리마다 렌더러를 손으로 구현한다. 메시지 포맷팅, 완성 파싱, 턴 간 연결을 모두 수작업으로 처리한다. 새 모델이 나올 때마다 새 렌더러가 필요하다. 현재 주요 오픈웨이트 패밀리를 커버하지만 유지보수 비용이 있다.

TITO 방식: “디코딩한 것은 다시 인코딩하지 않는다"는 규칙 하나로 설계한다. 템플릿이 prefix-preserving이면 — 거의 모든 템플릿이 그렇다 — 모델별 코드 없이 작동한다. 유지보수 부담이 훨씬 적다.

실무자가 볼 핵심 포인트

에이전트 RL을 구현하는 ML 엔지니어에게 이 문제는 디버깅이 극도로 어렵다. 코드가 정상 실행되고, 학습 곡선도 처음엔 그럴듯해 보인다. 그래디언트 신호가 조용히 손상돼 있을 뿐이다. 멀티턴 에이전트 학습에서 이상 현상이 보이면 재토크나이제이션 경로를 먼저 의심해야 한다.

실용적인 체크: 학습 루프에서 “메시지 리스트를 렌더링해서 토크나이즈"하는 부분이 롤아웃 중간 또는 끝에 있는지 확인한다. 있으면 TITO 원칙을 위반하고 있을 가능성이 높다. 실행 중 누적 토큰 버퍼로 교체하는 것이 해법이다.

prefix-preserving 검증: 새 모델에 TITO를 적용하기 전, 채팅 템플릿이 prefix-preserving인지 확인할 것. 검증법은 간단하다 — 동일한 메시지 prefix로 렌더링해 토큰이 바뀌지 않는지 확인한다.

단일 턴 RL에서 멀티턴 에이전트 RL로 넘어가는 팀이 많아지는 시점에, 이 원칙을 처음부터 설계에 반영하는 것이 나중에 디버깅하는 것보다 훨씬 저렴하다.

원문 출처

원문: Agentic RL: Token-In, Token-Out Done Right — Quentin Gallouédec, Kashif Rasul / HuggingFace Blog (2026-05-29)