확산 언어 모델의 추론 탐색을 넓히는 빈도 인식 리마스킹

확산 언어 모델(DLM)은 일반적인 자기회귀 모델처럼 텍스트를 한 방향으로 한 번에 생성하지 않는다. 여러 차례의 디노이징을 거치며 토큰을 확정하거나, 특정 위치를 다시 마스킹해 재예측한다. 어떤 토큰을 언제 확정하고 언제 되돌릴지는 최종 출력의 품질과 추론 비용을 모두 좌우한다. 너무 일찍 확정하면 약한 추론 경로에 고정될 수 있고, 반대로 지나치게 다시 마스킹하면 계산 자원이 낭비된다.

이번 연구는 DLM의 테스트 타임 스케일링을 다룬다. 기존의 확신도 기반 전략은 특정 시점의 확률이나 엔트로피를 사용해 토큰의 신뢰도를 판단한다. 하지만 수학처럼 어려운 추론 문제에서는 한 순간의 높은 확신도가 전체 추론의 정답을 보장하지 않는다. 모델이 그럴듯하지만 잘못된 경로를 일찍 선택하면, 추가 샘플링을 해도 모델이 선호하는 경로만 반복할 수 있다.

핵심 아이디어

• 2단계 계산 배분: 먼저 빠른 기본 샘플러로 모든 문제를 처리한다. 기본 전략으로 해결되지 않은 어려운 사례에만 추가 계산을 투입한다.
• 예측 변화 추적: 각 디노이징 단계에서 토큰의 예측이 몇 번 바뀌었는지 기록한다. 이는 현재의 확신도만 보는 방식과 달리, 모델이 해당 위치에서 얼마나 계속 탐색했는지를 보여주는 신호다.
• 빈도 인식 리마스킹: 예측이 자주 바뀐 토큰에 대해 리마스킹 방식을 조정한다. 모델이 가장 선호하는 하나의 추론 경로로 너무 빨리 수렴하지 않고, 다른 경로를 탐색하도록 유도하는 것이 목표다.
• 어려운 입력에만 적용: 모든 문제의 디코딩 비용을 높이는 대신, 기본 전략이 실패한 문제에 더 복잡한 전략을 집중한다.

실험은 LLaDA-8B-Instruct와 Dream-7B-Instruct를 대상으로 진행됐다. MATH-500에서 강력한 적응형 계산 배분 기준선과 비교했을 때, 절대 정확도 향상 폭은 LLaDA에서 최대 3.73%포인트, Dream에서 최대 7.33%포인트였다. HumanEval에서는 절대 커버리지 향상이 각각 최대 11.39%포인트와 10.98%포인트로 보고됐다. 또한 기본 확신도 기반 전략이 100회 샘플링 뒤에도 해결하지 못한 문제 가운데 3.75%에서 18.75%를 빈도 인식 리마스킹으로 추가 해결할 수 있었다.

의미와 한계

이 연구의 핵심은 단순히 샘플 수를 늘리는 데 있지 않다. 디노이징 과정에서 예측이 바뀐 이력을 순간적인 확신도의 보완 신호로 활용하고, 추가 테스트 타임 계산을 필요한 문제에 배분한다는 점이 중요하다. 반복적으로 중간 결정을 수정할 수 있는 DLM의 특성을 이용하면, 모델을 다시 학습하지 않고도 추론 탐색을 넓힐 가능성이 있다.

다만 제시된 결과는 두 모델과 논문에서 다룬 수학·코딩 벤치마크에 기반한다. 다른 작업에서도 예측 변화 빈도가 안정적인 신호로 작동하는지, 추가 리마스킹이 실제 서비스의 지연 시간과 계산 비용을 감수할 만큼 효과적인지는 더 검증해야 한다. 그럼에도 쉬운 문제는 빠르게 확정하고, 어려운 문제는 추론 경로를 더 오래 열어 두는 방식은 DLM 추론 설계에 유용한 방향을 제시한다.

출처: Hugging Face Daily Papers