[논문리뷰] Hierarchical Residuals: A Minimalist Decoupling of Residual Width from Compute Width

링크: 논문 PDF로 바로 열기

저자: Ziteng Wang, Jianqiao Lu, Xinyu Ma, Jintao Zhang, Jun Zhu, Yiyuan Ma, Jianfei Chen (Tsinghua University, ByteDance Seed)

1. Key Terms & Definitions (핵심 용어 및 정의)

• HiRes (Hierarchical Residuals): 인접한 Transformer layer들을 B개의 block으로 묶은 2단 Residual 구조입니다. Block 내부에서는 d-dimensional PreNorm inner residual이, block 사이에서는 D-dimensional PostNorm outer residual이 동작합니다.
• Residual Width vs. Compute Width: Residual width는 depth를 따라 전달되는 상태(state)의 차원입니다. Compute width는 각 attention/FFN sublayer가 입출력으로 처리하는 차원입니다. Vanilla Transformer에서는 두 값이 모두 d로 묶여 있습니다.
• Residual Expansion Ratio (n = D/d): Outer residual state가 inner compute width 대비 몇 배 넓은지를 나타내는 비율입니다. 기본 설정은 n = 2입니다.
• Read/Write Projection (Interface): Block 진입 시 D→d로 사상하는 dense read projection R_b와, 출력 시 d→D로 사상하는 dense write projection W_b입니다. Weight sharing 같은 structural inductive bias가 없는 가장 일반적인 linear interface입니다.
• HiNorm & Interface Stabilization: HiNorm은 D = d, R = W = I로 두고 outer residual을 제거한 control 구조로, L/B layer마다 RMSNorm만 삽입합니다. Interface stabilization은 learnable sigmoid gate로 identity(truncation Π)와 learned mapping을 interpolation하는 기법이며, 초기값은 σ(α) = 0.1입니다.

2. Motivation & Problem Statement (연구 배경 및 문제 정의)

본 논문은 Residual width와 Compute width의 결합을 풀되, 최소한의 아키텍처 변경과 standard operator만으로 넓은 residual state의 이점을 얻는 것을 목표로 합니다.

• 기존 Residual의 한계: Vanilla PreNorm Transformer는 depth를 따라 전달되는 state의 width를 각 layer의 working space width와 동일하게 강제합니다.
• 기존 확장 기법들의 공통점: AltUp, Hyper-Connections (HC), VWN, mHC, Attention Residuals 등은 두 가지 방향을 공유합니다. 하나는 effective residual state를 확장하는 것이고, 다른 하나는 residual accumulation의 scale을 명시적으로 regulate하는 것입니다.
• 기존 확장 기법들의 비용: 이 방법들은 per-layer 확장 state와 복잡한 mixing operator 때문에 activation traffic과 memory-bound 연산이 증가합니다. 그래서 throughput 회복을 위해 kernel fusion, recomputation, specialized kernel이 필요합니다.
• 핵심 질문: 이득이 특정 operator가 아니라 "넓은 state 자체"에서 온다면, 더 단순한 구조로도 같은 이득을 얻을 수 있는지를 묻습니다.

3. Method & Key Results (제안 방법론 및 핵심 결과)

HiRes는 Transformer layer 내부는 그대로 두고, block 경계에서만 dense read/write projection과 PostNorm을 추가해 residual width를 D로 확장합니다 [Figure 1].

Figure 1 — HiRes의 전체 아키텍처(inner/outer residual, read/write projection)를 보여주는 핵심 다이어그램

방법론

• 각 block b에서 x_{b,0} = R_b h_{b-1}로 outer state를 읽습니다.
• 이어서 2L/B개의 vanilla PreNorm sublayer를 d 차원에서 수행합니다.
• 마지막으로 h_b = RMSNorm(h_{b-1} + W_b Φ_b(x_{b,0}))로 outer residual에 기록합니다.
• Dense interface를 매 sublayer마다 쓰면 Θ(LDd) 비용이 발생하지만, block 경계에서만 사용해 Θ(BDd)로 amortize합니다.
• Outer PostNorm은 block마다 residual scale을 reset합니다. Inner PreNorm은 backward 시 optimization stability를 유지합니다.
• Embedding과 LM head에는 각각 projection E (d→D)와 H (D→d)를 추가합니다.
• 구현은 forward 함수 기준 10줄 미만의 코드 변경으로 가능합니다.

핵심 결과

• Model-size scaling (70M~626M, 5개 scale) [Figure 2]:
  • Fitted scaling law 기준으로 baseline은 validation loss 3.00에 도달하는 데 HiRes(n = 2) 대비 1.30×의 FLOPs가 필요합니다.
  • HiNorm(regulation)과 widening(n = 2)은 상호 보완적인 이득을 줍니다.
  • n = 1은 HiNorm과 비슷한 수준이고, n = 3, 4는 diminishing returns를 보입니다.
• Depth scaling (d = 512, 8→128 layers):
  • HiRes는 모든 depth에서 baseline보다 우수합니다.
  • L = 128에서 baseline의 validation loss는 2.9765, HiRes(n = 2, B = 8)는 2.9017입니다.
  • Baseline은 L = 64에서 128로 갈 때 scaling이 꺾이는 반면, B = 8 설정은 거의 직선 scaling을 유지합니다.
• 다른 확장 기법과의 비교 [Table 2]:
  • (768, 12) 설정에서 HiRes는 2× 확장만으로 validation loss 3.1300을 기록해 mHC-4(3.1459)보다 낮습니다.
  • Activation I/O는 4.58d로, mHC/VWN의 17d에 비해 크게 낮습니다.
  • Peak memory overhead는 3.3–3.4%, throughput은 baseline 대비 91.7–93.1%를 유지합니다.
  • (1280, 20) 설정에서는 평균 정확도 52.76%로 비교 기법 중 최고입니다.
  • Vanilla PyTorch 구현만으로도 Triton 구현과 거의 같은 throughput을 냅니다.
• OLMoE-1B-7B, 500B tokens:
  • Training loss가 0.0208 낮아졌습니다 (2.3663 → 2.3455).
  • Validation loss는 0.0180 낮아졌습니다.
  • 12개 downstream task의 평균 정확도는 1.14 point 향상되었습니다 (64.24 → 65.38).
  • 추가 파라미터는 83.9M으로, active parameter의 6.5%입니다.
• Ablation:
  • Interface를 frozen random으로 고정하면 validation loss가 0.031 증가합니다.
  • Boundary RMSNorm을 제거하면 0.025 증가합니다.
  • Outer residual을 제거하면 0.032 증가합니다.
  • n = 0.75에서는 HiNorm보다도 개선되지 않았습니다. 이는 이득의 주원인이 interface 자체가 아니라 residual width임을 뒷받침합니다.

Figure 2 — 1.30× FLOPs efficiency를 보여주는 핵심 scaling 결과

Table 2 — 다른 residual expansion 기법 대비 품질과 효율을 종합 비교한 핵심 테이블

4. Conclusion & Impact (결론 및 시사점)

HiRes는 dense projection, normalization, residual addition만으로 residual width를 compute width로부터 분리할 수 있음을 보였습니다. 이를 통해 두 width가 같아야 한다는 관례가 suboptimal일 수 있음을 실증했습니다.

• 성능 측면: 모델 크기, depth, 대규모 MoE pretraining 전반에서 일관된 개선을 보였습니다. HC, mHC, VWN, AttnRes 대비로도 경쟁력 있는 품질을 더 낮은 memory와 I/O 비용으로 달성합니다.
• 산업적 의미: Specialized kernel이나 infrastructure-level 최적화 없이 기존 학습 스택에 바로 적용할 수 있습니다.
• 학술적 의미: Residual width를 독립적인 scaling variable로 다룰 수 있는 실용적 근거를 제시합니다. 또한 MLA, LatentMoE 같은 latent-space 설계와 연결되는 해석을 제공합니다.
• 한계: 다음 사항들은 향후 과제로 남아 있습니다.
  • Fixed compute budget 하에서 width와 depth의 공동 최적 배분은 탐구되지 않았습니다.
  • Non-uniform block 구성은 검토되지 않았습니다.
  • 넓은 state가 어떤 정보를 보존하는지에 대한 메커니즘 분석이 부족합니다.

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.