플립플롭을 멈춰라 빠른 복구 가능한 확산 디코딩을 위한 컨텍스트 보존 검증

초록

본 논문은 확산 언어 모델의 병렬 디코딩 과정에서 발생하는 불필요한 재마스킹과 토큰 복구를 ‘플립플롭’이라 정의하고, 이를 최소화하기 위해 KV 캐시를 오버라이드하는 단일 패스 검증 기법인 COVER를 제안한다. 시드 토큰을 선택적으로 마스킹하면서 이전 단계의 KV 상태를 재사용해 컨텍스트 손실을 방지하고, 대각선 보정으로 자기 누수를 차단한다. 안정성 기반 시드 선택과 적응형 검증 수 조절을 통해 검증 비용을 크게 줄이며, 기존 방법 대비 11배 가량의 속도 향상을 달성한다.

상세 분석

COVER의 핵심 아이디어는 두 가지 상반된 요구를 하나의 포워드 패스에 결합하는 것이다. 첫째, 시드 토큰을 ‘leave‑one‑out’ 방식으로 검증하려면 해당 위치를 마스크해야 하지만, 마스크는 다른 토큰들의 어텐션 컨텍스트를 약화시킨다. 둘째, 병렬 초안(drafting) 단계에서는 시드 토큰의 기존 표현이 그대로 유지돼야 효율적인 토큰 예측이 가능하다. 이를 해결하기 위해 저자들은 KV 캐시 오버라이드 메커니즘을 설계하였다. 구체적으로, 시드 위치를 마스크한 입력을 만든 뒤, 해당 위치의 키·밸류(K,V) 쌍만 이전 단계에서 저장한 캐시값으로 교체한다. 이렇게 하면 비시드 쿼리는 여전히 원래 시드 토큰의 정보를 참조할 수 있어 어텐션 컨텍스트가 유지된다.

하지만 시드 자체를 검증하기 위해서는 자기 자신에 대한 직접적인 키·밸류 연결을 차단해야 한다. 저자들은 대각선 보정(diagonal correction)이라는 수식을 도입해, 시드 위치의 자기 어텐션 점수를 원래 마스크된 상태의 값으로 복원하고, 전체 어텐션 분포가 소프트맥스 정규화 과정에서 왜곡되지 않도록 스케일링 계수를 적용한다. 이 과정은 수학적으로 닫힌 형태로 표현되며, 구현 시 추가 연산 비용이 거의 없다는 장점이 있다.

시드 선택 전략 역시 중요한 기여점이다. 기존 방법은 단순히 확신도가 낮은 토큰을 재검증했지만, 이는 종종 ‘플립플롭’ 현상을 유발한다. COVER는 불확실성(엔트로피), downstream 영향력(시드가 다른 마스크 토큰에 미치는 어텐션 기여), 그리고 캐시 드리프트(시간에 따른 KV 변화)를 종합한 안정성 점수를 정의한다. 이 점수는 토큰이 현재 단계에서 얼마나 위험한지와 동시에 캐시 재사용이 얼마나 신뢰할 수 있는지를 동시에 평가한다. 또한, 매 단계마다 검증할 시드 수를 동적으로 조절해, 예산을 효율적으로 배분한다.

실험 결과는 세 가지 주요 지표에서 기존 방법을 앞선다. 첫째, 플립플롭 횟수가 90% 이상 감소한다. 둘째, 전체 디코딩 단계 수가 평균 30%~70% 줄어들어 실제 추론 속도가 2×~11× 가속된다. 셋째, BLEU, ROUGE, GPT‑4 평가 등 품질 지표에서는 기존 방법과 동등하거나 약간 개선된 결과를 보인다. 특히 대규모 모델(Dream‑Instruct‑7B, LLaDA‑Instruct‑8B)에서 효과가 두드러지며, 메모리 사용량도 KV 오버라이드만으로 기존 방법보다 낮다.

종합하면, COVER는 ‘검증 = 재마스킹’이라는 전통적 패러다임을 깨고, 컨텍스트 보존과 정확한 leave‑one‑out 검증을 동시에 달성함으로써 확산 언어 모델의 병렬 디코딩을 실용적인 수준으로 끌어올린 혁신적인 접근이다.