앤트 UBTree 추론 디코딩 개선: 초안 트리 병렬 필터링으로 대형 모델이 한 번에 더 많은 토큰 검증 가능

동찰 Beating AI 속보, 앤트백링이 대형 모델 추측 디코딩 방안 UBTree를 오픈소스로 공개했다.

추측 디코딩은 소형 모델이 미리 '초안'을 작성하고, 대형 모델이 여러 Token을 한 번에 검증하여 Token을 하나씩 생성하는 데 필요한 계산 라운드를 줄이는 방식이다.

그러나 초안에도 적중률 문제가 있다. 소형 모델이 일찍 틀릴수록 이후에 직접 채택할 수 있는 내용이 줄어든다. 이전에도 여러 후보 경로를 동시에 준비하여 초안 트리로 구성한 뒤, 수용될 가능성이 낮은 분기를 가지치기하는 기술이 있었다.

UBTree는 선별 방법을 개선했다. 먼저 이후 몇 개 위치에 나타날 수 있는 Token을 병렬로 예측한 다음, 경량 모델로 인접 Token이 합리적으로 연결될 수 있는지 판단한다. 서로 다른 조합을 동시에 평가할 수 있어 각 경로를 따라 모델을 반복 실행할 필요가 없다. 이후 더 유망한 분기를 골라 초안 트리로 구성하여 대형 모델에 검증을 맡긴다.

팀은 선별 모델의 훈련 방식도 조정하여 확률이 가장 높은 답변만 주목하는 것이 아니라 유지할 가치가 있는 대안 경로도 식별할 수 있게 했으며, 이를 통해 대형 모델이 매 라운드 더 많은 Token을 수용할 기회를 갖게 했다.

Qwen3-4B의 7개 테스트에서 DFlash 2 선택기로 트리를 구성한 평균 가속비는 5.95배였다. UBTree 선택기로 교체하고 훈련 방법을 조정한 후 6.88배로 향상되어, 기존 방안 대비 약 16% 추가 개선되었다.