Wavefront Decoding, 루프형 언어 모델의 디코딩을 병렬화하다

서론: 반복 구조가 만드는 지연

루프형 언어 모델은 동일한 가중치 공유 블록을 여러 번 적용해 파라미터 수를 크게 늘리지 않고도 유효 계산 깊이를 확장합니다. 그러나 생성 과정에서는 토큰 하나를 만들 때마다 반복 블록을 여러 차례 순서대로 실행해야 할 수 있습니다. 가중치 공유가 모델 규모 측면에서는 장점이지만, 디코딩 지연이라는 새로운 병목을 만드는 셈입니다.

Wavefront Decoding(WFD)은 모델을 다시 학습하거나 별도의 초안 모델을 추가하는 대신, 루프형 모델 내부의 계산 순서를 재구성해 이 문제를 해결하려 합니다. 핵심은 중간 반복 상태를 활용하고, 서로 다른 토큰 위치와 반복 깊이의 상태를 하나의 배치 호출에서 처리하는 것입니다.

핵심 아이디어: 초안 생성과 검증을 겹치기

일반적인 추측 디코딩은 보통 두 단계로 구성됩니다. 먼저 빠른 초안 모델이 후보를 만들고, 이후 대상 모델이 후보를 검증합니다. WFD는 루프형 아키텍처의 특성을 이용해 이 단계를 분리하지 않습니다.

• 중간 출력으로 초안을 만든다. 토큰 상태가 전체 반복 깊이에 도달하지 않았더라도 다음 토큰 예측을 만들 수 있으므로, 얕은 단계의 상태를 새 위치의 초안 생성에 활용합니다.
• 가중치 공유를 배치 처리에 사용한다. 위치와 반복 깊이가 서로 다른 상태도 같은 블록을 통과하므로, 하나의 반복 블록 호출에서 함께 처리할 수 있습니다.
• 대각선 파형으로 작업을 배치한다. 새 위치는 얕은 깊이에서 시작하고, 앞서 생성된 위치는 더 깊은 반복으로 이동해 전체 깊이 검증을 받습니다.
• 거부된 초안을 보정한다. 얕은 예측이 검증을 통과하지 못하면 전체 깊이에서 얻은 예측을 사용해 결과를 수정합니다.

따라서 WFD의 핵심은 단순히 배치 크기를 키우는 데 있지 않습니다. 토큰 위치와 반복 깊이에 걸쳐 있던 의존성을 파형 형태로 재배열해, 초안 생성과 검증을 연속적인 반복 호출 안에서 동시에 진행하는 데 있습니다.

보고된 결과와 KV 최적화

논문 초록에 따르면 WFD는 6개 Spec-Bench 작업 범주에서 평가되었습니다. 자기회귀 디코딩과 비교했을 때 Ouro-2.6B에서는 2.42배, Huginn-3.5B에서는 3.54배의 속도 향상을 기록했습니다. 또한 초안을 모두 만든 뒤 검증하는 단계 분리 방식보다 일관되게 높은 성능을 보였다고 보고합니다.

저자들은 파형 처리에서 발생하는 KV 트래픽을 줄이기 위해 반복 단계 간 KV 공유도 제안했습니다. 이 최적화를 적용하면 Huginn-3.5B에서 보고된 속도 향상은 최대 4.81배로 증가합니다. 다만 수치는 특정 모델, 벤치마크, 구현 조건에서 얻은 결과이므로 모든 루프형 모델이나 하드웨어에 동일하게 적용된다고 볼 수는 없습니다. 관련 코드는 공개되어 후속 검증이 가능합니다.

의미와 영향

WFD는 루프형 모델의 반복 계산을 단순한 직렬 비용이 아니라 조정 가능한 추론 자원으로 바라봅니다. 중간 상태의 예측력을 활용하고, 가중치 공유를 통해 서로 다른 깊이의 상태를 묶어 처리함으로써 별도의 초안 모델을 학습하지 않고도 디코딩을 가속할 가능성을 제시합니다.

물론 실제 이득은 중간 예측의 품질, 파형 스케줄러 구현 난도, KV 캐시의 메모리 대역폭에 영향을 받습니다. 그럼에도 이 연구는 반복형 아키텍처에서 모델 깊이 방향의 계산을 어떻게 추론 병렬성으로 전환할지라는 중요한 시스템 설계 문제를 구체적으로 보여줍니다.

출처: Hugging Face Daily Papers