기술 요약: Dual-Flow Transformer

문제 정의

대규모 언어 모델(LLM)을 서빙하는 비용은 학습보다 추론 단계에서 점점 더 지배적으로 높아지고 있습니다. 그러나 추론의 두 가지 주요 단계인 프리필(prefill)(프롬프트 처리)과 자기회귀 디코딩(autoregressive decode)(토큰 생성)은 근본적으로 다른 하드웨어 병목 현상을 보입니다:

• 프리필: 프롬프트 토큰의 병렬 처리는 일반적으로 산술 처리량에 의해 제한되는 연산 집약적(compute-bound) 특성을 가집니다.
• 디코딩: 순차적인 생성은 각 새로운 토큰마다 성장하는 키-값(KV) 캐시와 모델 가중치를 반복해서 읽어야 하므로, 일반적으로 메모리 대역폭 제한적(memory-bandwidth-bound) 특성을 가집니다.

기존의 스케일링 전략(모델 너비나 깊이 증가)은 모든 레이어가 프리필과 디코딩 단계 모두에서 평가되기 때문에 두 단계의 비용을 동시에 높입니다. 이는 모델의 품질을 개선하기 위해 추가적인 연산을 도입하면 필연적으로 프롬프트 처리 비용까지 증가하게 만드는 경정성을 만듭니다. 본 논문은 단일한 지속적 KV 캐시와 프롬프트 전반의 연산을 유지하면서도, 오직 연속 예측(continuation prediction)(디코딩)에만 추가적인 학습된 연산을 할당할 수 있는 방법이 있는지에 대해 질문합니다.

방법론: Dual-Flow Transformer

저자들은 프라이머리 프롬프트 처리 경로와 보조 연속 예측 경로를 구조적으로 분리하는 Dual-Flow Transformer 아키텍처를 제안합니다.

핵심 아키텍처

이 모델은 모델의 대부분의 파라미터를 공유하지만 상태 관리와 활성화 범위가 서로 다른 두 개의 잔차 상태 궤적(residual state trajectories)으로 구성됩니다:

1. Primary Flow (주 경로): 완전하고 표준적인 인과적(causal) 트랜스포머 경로입니다. 전체 프롬프트를 처리하고, 지속적인 KV 캐시를 기록하며, 주요 상태를 결정합니다. 이는 자기 완결적이며 보조 흐름에 의존하지 않습니다.
2. Auxiliary Flow (보조 경로): 주 경로의 상태와 캐시를 읽지만, 지속적인 KV 캐시에 절대 쓰지 않으며 주 경로로 피드백되지도 않는 잔차 궤적입니다.
   • 활성화: 보조 흐르는 초기 프롬프트 처리(prefill) 단계에서는 생략됩니다. 마지막 프롬프트 위치(첫 번째 연속 토큰을 예측하기 위함)와 그 이후의 모든 생성 위치에서 활성화됩니다.
   • 가중치 공유: 두 흐름은 모든 주요 어텐션, MLP, 출력 행렬을 공유합니다.
   • 분리: 이들은 별도의 토큰 임베딩과 경량화된 학습된 결합 벡터(coupling vectors)를 사용합니다.

핵심 메커니즘

• 비대칭 공유-KV 어텐션 (Asymmetric Shared-KV Attention): 주 경로는 표준적인 인과적 셀프 어텐션을 수행합니다. 보조 흐름은 새로운 쿼리($Q_2$)를 생성하며, 이는 학습된 벡터($a_\ell$)를 통해 주 쿼리($Q_1$)와 결합됩니다. 결정적으로, 보조 흐름은 주 경로의 키(Keys)와 값(Values)($K_1, V_1$)에 어텐션을 수행합니다. 이를 통해 보조 흐름은 자신의 캐시를 쓰지 않고도 주 경로의 캐시된 상태를 활용할 수 있으며, 가중치와 KV 영역이 한 번 로드되어 두 흐름 모두에 서비스되는 그룹화된 실행을 가능하게 합니다.
• 경량 결합 (Lightweight Coupling): 정보는 쿼리, MLP 중간층, MLP 출력 레벨에서의 학습된 벡터($a_\ell, b_\ell, c_\ell$)를 통해 주 경로에서 보조 경로로 단방향으로 흐릅니다. 이는 주 경로의 경로를 변경하는 의존성을 만들지 않으면서 보조 표현을 풍부하게 만듭니다.
• 혼합 목적 함수 (Mixture Objective): 최종 예측은 주 경로와 보조 경로의 분포의 혼합입니다: $p_{mix} = \alpha_t p_{primary} + (1-\alpha_t) p_{aux}$. 혼합 가중치 $\alpha_t$는 주 경로의 신뢰도(충돌 확률)에 따라 동적으로 조정되며, 추론 중에 주 경로가 지배적(최소 50% 가중치)으로 유지되도록 보장합니다.
• 라우터 리플레이 (Router Replay, MoE용): 혼합 전문가(Mixture-of-Experts, MoE) 구성에서, 보조 흐름은 주 경로가 결정한 동일한 전문가 집합을 재사용하지만, 이를 자신의 혼합 가중치와 함께 적용합니다. 이는 주 경로의 전문가 선택 로직을 보존하면서 전문가에 수행되는 산술 연산을 증가시킵령합니다.

추론 vs. 학습 비용

• 추론: 프리필 단계에서는 추가 비용이 발생하지 않습니다. 보조 흐로는 완전히 건너뛰어집니다. 디코딩 단계에서는 추가적인 연산(토큰당 산술 연산이 대략 두 배)이 발생하지만, 두 흐름이 그룹화되어 있으므로 KV 캐시 크기나 고유 가중치 세트가 늘어나지는 않습니다.
• 학습: 두 흐름 모두 전체 시퀀스에 대해 평가되므로, 주 모델 대비 약 2배의 학습 FLOPs가 소요됩니다. 그러나 파라미터 수의 증가는 미미합니다(별도의 임베딩과 결합 벡터만 추가됨).

주요 기여

1. 아키텍처: 프롬프트 전반의 계산과 연속 계산을 분리하는 비대칭 공유-KV 설계를 도입했습니다. 주 경로는 상태 지속성을 담당하고, 보-조 경로는 예측력을 높이기 위한 "잔차" 궤적을 제공합니다.
2. 경험적 품질: 세 가지 실험 축에서 검증 손실(validation loss)의 일관된 감소를 입증했습니다:
   • 데이터 스케일링 (NanoGPT): Dual-Flow는 동일한 토큰 예산에서 표준 트랜스포머보다 우수한 성능을 보이며, 더 낮은 적합 손실 하한선(fitted loss floor)을 가집니다.
   • 모델 크기 스케일링 (Dense LLaMA): 모델 크기가 커짐에 따라 이 이점은 지속됩니다.
   • 희소 MoE: 라우터 리플레이를 포함한 희소 LLaMA 스타일의 MoE 구성에서도 유효성이 검증되었습니다.
3. 단계별 MoE 할당: 기존의 병렬 스트림 설계가 균일한 연산 할당을 적용하는 것과 달리, Dual-Flow는 주 경로와 보조 경로의 전문가 팬아웃(fan-out)을 독립적으로 구성할 수 있게 합니다. 이를 통해 프리필 비용, 디코딩 비용, 예측 품질 사이의 트레이드오프 공간을 제공하여 특정 워크로드에 최적화된 시스템을 구축할 수 있습니다.

결과

• 검증 손실: NanoGPT 실험의 매칭된 토큰 예산 비교에서, Dual-Flow는 표준 트랜스포머보다 일관되게 낮은 검증 손실을 달s했습니다. 적합된 점근적 손실 하한선은 2.9416(Transformer)에서 2.9013(Dual-Flow)으로 감소했습니다.
• 연산 효율성: 학습-연산 매칭 비교(Dual-Flow 모델이 절반의 데이터 스텝을 사용하여 Transformer가 두 배의 스텝을 수행할 때와 거의 동일한 공유 백본 FLOPs를 수행하는 경우)에서, Dual-Flow는 여전히 더 낮은 검증 손실을 달성했습니다. 이는 상호작용하는 흐름 사이에 연산을 할당하는 것이 단일 흐름을 위해 데이터를 단순히 늘리는 것보다 더 효과적임을 시사합니다.
• 절제 연구 (Ablations): 구성 요소 분석 결과, 최소한의 Dual 구조(별도 임베딩, 공유 KV)만으로도 가장 유사한 선행 연구(PHD-2)보다 우수한 성능을 보였습니다. 교차 흐름 결합(cross-flow coupling)과 확률 혼합을 추가했을 때 성능이 더욱 향상되었습니다.

의의 및 주장

본 논문은 Dual-Flow가 단계별 전문가 할당을 위한 구조적 기회를 확립한다고 주장합니다. 프롬프트 경로와 연속 경로를 분리함으로써, 이 아키텍처는 서빙 시스템이 프롬프트 처리 비용을 비례적으로 높이지 않고도 생성 시점의 연산을 확장 가능한 자원으로 취급할 수 있게 합니다.

특히 저자들은 MoE 모델에서 주 경로와 보조 경로가 독립적인 전문가 팬아웃을 가질 수 있다는 점을 강조합니다. 이는 세 가지 방향의 트레이드오프 노브(knob)를 생성합니다:

1. 프리필 비용: 주 경로의 전문가 사용량에 의해 결정됩니다.
2. 디코딩 비용: 주 경로와 보조 경로의 전문가 사용량 합계에 의해 결정됩니다.
3. 예측 품질: 연속 단계에 할당된 총 연산을 늘림으로써 향상됩니다.

이를 통해 시스템은 공유 파라미터 풀이나 단일 주 경로 캐시 구조를 변경하지 않고도, 프롬프트 계산 비용은 유지하면서 가용한 디코딩 예산을 더 많은 보조 전문가에 사용하거나, 그 반대로 설정할 수 있는 배포 특화형 유연성을 갖게 됩니다. 이 연구는 다음 토큰 계산이 외부적인 긴 추론 궤적이나 다중 샘플링을 통해서가 아니라, 모델 내부에서 할당될 수 있는 유용한 스케일링 자원이라는 전제를 뒷받침합니다.