기술 요약: WHALE - Wukong-HSTU 아키텍처를 통한 확장 가능한 통합 추천 모델

1. 문제 정의

대규모 추천 시스템은 두 가지 구별되면서도 상호 보완적인 랭킹 신호 유형을 모델링해야 합니다: 비순차적 특징(사용자/아이템 속성, 컨텍스트 및 명시적 교차 특징)과 순차적 특징(사용자 행동 이력).

• 현재의 한계: 기존 아키텍처들은 일반적으로 이러한 신호들을 별개로 취급합니다. Wukong 아키텍처는 비순차적 특징 간의 고차 상호작용을 확장하는 데 탁월하지만, 행동 시퀀스의 시간적 진화를 모델링하는 메커니즘이 부족합니다. 반대로, HSTU(Hierarchical Sequence Transduction Unit) 아키텍처는 길고 이질적인 행동 시퀀스를 효과적으로 모델링하지만, 사용자, 아이템, 요청 특징 간의 고차 정적 및 컨텍스트 교차를 명시적으로 확장하지는 못합니다.
• 공백(The Gap): 행동 이력을 고정된 요약본으로 압축하여 특징 상호작용 모델에 입력하는 얕은 결합(shallow-hybrid) 방식은 병목 현상을 일으킵니다. 일단 압축되면, 고차 특징 상호작용은 사용자의 이력 중 특정 이벤트로부터 미세한 증거를 선택적으로 검색할 수 없게 됩니다. 이는 정적 특징 상호작용과 동적 순차 패턴 사이의 "점진적 교환(progressive exchange)"을 제한하며, 서로 다른 후보군에 대해 동일한 이력의 서로 다른 부분이 필요할 때 발생하는 관련성 문제를 해결하는 데 한계가 있습니다.

2. 방법론: WHALE 아키텍처

저자들은 Wukong과 HSTU를 단일 모델 내에서 통합하고 두 백본을 네트워크 전체에서 활성화 상태로 유지하는 재귀적 아키텍처인 WHALE(Wukong-HSTU-based scALable unified modEl)을 제안합니다.

2.1 핵심 아키텍처

WHALE는 입력 레이어와 그 뒤를 잇는 재귀적 WHALE 레이어 스택으로 구성됩니다. 각 레이어는 세 개의 뚜렷한 모듈을 포함합니다:

1. Wukong 모듈: 비순차적 브랜치(사용자, 아이템, 컨텍스트 및 교차 특징)에서 작동합니다. 스택형 Factorization-Machine 블록(FMB)과 선형 압축 블록(LCB)을 사용하여 고차 상호작용 표현($H_w$)을 구축합니다.
2. HSTU 모듈: 순차적 브랜치(사용자 행동 이력)에서 작동합니다. 상대적 어텐션 바이어스(relative attention bias)를 갖춘 어텐션 메커니즘을 사용하여 미세한 시간적 증거와 의존성을 인코딩하고, 전체 시퀀스 길이($L$)를 유지하는 시퀀스 인식 행동 표현($H_s$)을 생성합니다.
3. 어텐션 기반 퓨전 모듈: 점진적 교환을 가능하게 하는 핵심 혁신입니다.
   • 메커니즘: Wukong 상호작용 표현($H_w$)을 **쿼리(Query)**로, HSTU 행동 표현($H_s$)을 **키(Key) 및 값(Value)**으로 취급합니다.
   • 기능: 이를 통해 후보 및 컨텍스트에 종속적인 특징 교차가 사용자의 이력으로부터 가장 관련성이 높은 행동 증거를 선택적으로 검색할 수 있게 합니다.
   • 정교화: 어텐션 출력은 원래의 Wukong 표현과 결합되고, Fusion MLP를 통해 투영된 후, 프리-노름(pre-norm) SwiGLU FFN(Gated Feed-Forward Network)을 통과하여 융합된 표현($H_f$)을 생성합니다.
   • 전파: 융합된 표현($H_f$)은 다음 레이어의 비순차적 입력이 되는 반면, HSTU 출력($H_s$)은 다음 레이어의 시퀀스 모듈으로 독립적으로 전파됩니다. 이 설계는 정적 상호작용과 동적 패턴이 여러 레이어에 걸쳐 반복적으로 서로에게 조건화(condition)되도록 보장합니다.

2.2 효율성 최적화

WHALE를 산업 현장에 배포 가능하게 만들기 위해 저자들은 다음과 같은 시스템 레벨의 최적화를 구현했습니다:

• 맞춤형 Triton 커널:
  • 공유 키-값(Shared Key-Value): 교차 브랜치 어텐션은 메모리 대역폭에 민감하므로, 저자들은 Key와 Value 텐서($K=V$)를 공유하여 GPU 메모리 트래픽을 절반으로 줄였습니다.
  • 비대칭 역전파 변형(Asymmetric Backward Variants): 비대칭적인 형태(적은 수의 쿼리 vs 긴 이력)를 인식하여, 그래디언트 감소 비용을 최소화하기 위해 KV-병렬 또는 Q-병렬 역전파 스케줄 사이를 동적으로 선택합니다.
• 공유 게이트 SwiGLU FFN: SwiGLU FFN의 게이트 및 업-프로젝션 행렬을 결합($W_{gate} = W_{up}$)하여 품질 저하 없이 행렬 곱셈을 3개에서 2개로 줄였습니다.
• 혼합 정밀도 및 컴파일러 최적화: 밀집(dense) 백본은 BF16으로 실행되는 반면, 민감한 구성 요소(손실 함수, 태스크 아키텍처)는 FP32로 유지됩니다. 연산자를 퓨전하고 메모리 사용량을 최적화하기 위해 torch.compile이 사용되었습니다.
• 추론 최적화: 연산자 퓨전을 위한 AOTInductor와 동적 형태 계산 중 GPU-CPU 동기화 병목을 제거하기 위한 "형태 힌트 텐서(shape hint tensors)"를 포함하여 지연 시간을 1~5ms 줄였습니다.

3. 주요 기여

1. 통합 아키텍처: WHALE는 얕은 결합 대신 레이어별 점진적 교차 브랜치 교환을 활용하여 Wukong(비순차적 특징 상호작용)과 HSTU(순차적 행동 모델링)를 단일 모델 내에서 통합한 최초의 확장 가능한 추천 아키텍처로 제시됩니다.
2. 시스템-모델 공동 설계: 엄격한 지연 시간 및 처리량 제약 조건 하에서 이러한 복잡한 통합 모델의 온라인 배포를 가능하게 하는 맞춤형 커널 및 훈련/추론 최적화 개발.
3. 경험적 검증: 대규모 숏폼 비디오 소셜 미디어 플랫폼(80B 훈련 샘플, 4B 평가 샘플)에서 WHALE가 시퀀스 길이, 모델 깊이, 모델 너비에 따라 효과적으로 확장되며, 오프라인 및 온라인 설정 모두에서 강력한 단일 패러다임 베이스라인을 능가함을 입증했습니다.

4. 실험 결과

실험은 대규모 숏폼 비디오 소셜 미디어 플랫폼에서 수행되었습니다.

• 오프라인 성능 (확장성):

  • 아키텍처 비교: WHALE는 다양한 모델 복잡도(8, 14, 32 GFLOPs)에 걸쳐 Wukong 전용 및 HSTU 전용 베이스라인보다 일관되게 우수한 성능을 보였습니다.
  • 스케일링 법칙(Scaling Laws): WHALE는 시퀀스 길이(최대 15k), 모델 깊이(최대 8 레이어), 모델 너비(최대 512 차원)가 증가함에 따라 일관된 이득을 보여주었습니다.
  • 절제 연구(Ablation Studies): 어텐션 메커니즘을 제거(평균 풀링으로 대체)하거나 Fusion MLP/SwiGLU FFN을 제거하면 품질이 크게 저하되어, 특정 퓨전 설계의 필요성을 확인했습니다.

• 온라인 성능 (A/B 테스트):

  • 14일 동안 프로덕션에 배포되었습니다.
  • 주요 지표: 온라인 베이스라인 대비 **+0.113%**의 이득을 달了.
  • 보조 지표: 보조 지표에서 각각 +0.824% 및 +1.820%의 이득을 기록했습니다.
  • 트레이드오프: 모델은 **-5%**의 추론 QPS(처리량) 저하를 초래했으나, 저자들은 이것이 서빙 예산 범위 내에서 수용 가능한 수준이라고 언급했습니다.

5. 의의 및 주장

본 논문은 WHALE가 두 가지 구별되는 정보원인 정적 특징 상호작용과 동적 순차 행동을 산업 환경에서 어떻게 확장 가능하게 통합할 수 있는지에 대한 실질적인 사례를 제공한다고 주장합니다.

저자들은 산업용 추천 시스템의 다음 단계 확장이 단순히 개별 백본을 더 크게 만드는 것에서 오는 것이 아니라, 각 레이어에서 보완적인 아키텍처를 점진적으로 통합하는 데서 올 것이라고 주장합니다. 고차 특징 교차가 긴 사용자 이력으로부터 미세한 증거를 반복적으로 검색할 수 있도록 함으로써, WHALE는 얕은 결합의 한계를 해결하고 심층적이고 교차적인 브랜치 상호작용이 대규모 랭킹 시스템을 위한 유효하고 효과적인 설계 원칙임을 입증합니다.