Diffusion Language Model(DLM)은 Autoregressive Language Model(AR LM)과 달리, 여러 위치의 masked token을 반복적으로 예측하고 복원하면서 문장을 생성합니다. 이러한 방식은 양방향 문맥을 활용할 수 있을 뿐 아니라, 여러 token을 한 번에 생성할 수 있어 기존의 순차적 생성 방식보다 빠른 decoding이 가능하다는 장점이 있습니다. 하지만 여러 token을 동시에 생성할 수 있다는 것이 곧 효율적인 decoding을 의미하지는 않습니다. 실제 생성 속도는 각 step에서 어떤 token을 먼저 확정하고, 남아 있는 masked token을 어떤 순서로 복원하는지에 크게 영향을 받습니다.
기존의 parallel decoding 방법들은 주로 모델의 prediction confidence를 활용합니다. 예측 확률이 높은 token을 우선적으로 commit함으로써 생성 품질을 유지하면서 decoding step을 줄이는 방식입니다. 그러나 confidence가 높은 token을 먼저 선택하는 것이 항상 효율적인 decoding 순서를 만드는 것은 아닙니다. 예를 들어, 연속된 masked interval에서 한쪽 끝에 가까운 token을 먼저 선택하면, 해당 token을 commit한 이후에도 긴 masked interval이 그대로 남게 됩니다. 반면 중앙에 위치한 token을 선택하면 남은 interval이 두 개의 작은 구간으로 나뉘고, 이 구간들을 병렬적으로 처리할 수 있습니다.
즉, 어떤 token을 선택하느냐뿐 아니라, 그 token이 어디에 위치하는지도 전체 decoding 효율을 결정하는 중요한 요소입니다.
위의 그림에서 볼 수 있듯이, boundary에 가까운 pivot을 반복적으로 선택하면 긴 masked interval이 남아 decoding이 느리게 진행됩니다. 반면 interval의 중앙에 가까운 pivot을 선택하면 남은 구간들이 균형 있게 분할되어 병렬적으로 복원됩니다. 이상적인 balanced pivot selection에서는 길이가 J인 interval을 복원하는 데 필요한 decoding round를 O(J)에서 O(log(J+1)) 수준으로 줄일 수 있습니다. 따라서 우리는 prediction confidence뿐 아니라 token의 위치를 함께 고려하여, 생성 품질과 decoding 효율을 동시에 개선할 수 있는 새로운 방법을 제안합니다.
우리는 이러한 관찰을 바탕으로, token의 prediction confidence와 positional balance를 함께 고려하는 Balance-aware Pivot Decoding (BPD)를 제안합니다.
<aside> 💡
핵심아이디어
Confidence가 높은 token을 먼저 확정하고, 남은 masked interval에서는 가능한 한 중앙에 가까운 token을 선택하여 균형 있는 분할을 유도함
</aside>
BPD는 크게 두 단계로 구성됩니다.
또한 decoding 과정에서 Prefix KV Caching을 활용하여 반복적인 연산 비용을 줄입니다.