HRM-Text: Scaling을 넘어서는 효율적 프리트레이닝

HRM-Text는 듀얼 타임스케일 재귀와 MagicNorm, warmup deep credit assignment를 통해 긴 순환에서의 학습 안정성을 확보한다. 또한 task-completion objective와 PrefixLM masking을 결합해 Instruction-Response 데이터로부터 효율적으로 학습하며, 소규모 예산으로도 오픈형 모델과 경쟁 가능한 성능을 달성한다. 이로써 아키텍처와 학습 목표의 공동 설계가 compute-대-성능 비율을 크게 낮출 수 있음을 시사한다.

출발점: 현재의 대형 언어 모델 전재가 대규모 데이터/계산에 의존하는 한계가 있다. Transformer의 깊은 구조에서 gradient가 소실되거나 과도하게 커지는 문제를 재귀적으로 해결하는 방법은 존재하나, 긴 종속성을 다루기엔 불안정하다. 해결 원리: HRM은 L 모듈(빠른 업데이트)과 H 모듈(느린 맥락 유지)의 듀얼 타임스케일 재귀로 내부 계산을 다층적으로 처리하고, MagicNorm와 warmup deep credit assignment로 TBPTT 하에서의 활성화 분산과 그래디언트 흐름을 안정화한다. 목표 및 입력 구조의 변화: Broad raw-text pretraining 대신 instruction-response 쌍에 대해 −log P(xa|xq)를 학습하고 PrefixLM 마스크를 적용해 인스트럭션 토큰은 양방향 주의를 허용하되 응답 토큰은 자동회귀로 생성한다. 결과적으로 소량의 데이터와 토큰으로도, HRM-Text는 1B 파라미터에서 40B 토큰의 예산으로 대형 baselines에 근접한 성능을 달성하며, 학습 효율성 측면에서 큰 폭으로 개선된다. 달라지는 점: HRM은 깊은 순환의 효과적 깊이(effective depth)를 증가시키고, 실험적으로 gradient 신호를 유지하는 능력과 출력 분포의 안정화 여부를 개선한다. 이를 통해 아키텍처-목적 연합이 compute-to-performance를 크게 낮출 수 있음을 확인한다.

1. 전체 접근: L 모듈(빠른 업데이트)과 H 모듈(느린 업데이트)로 구성된 HRM-Text를 구성하고, 매 사이클마다 3회 L 모듈 업데이트 + 1회 H 모듈 업데이트를 수행하며, 마지막 H 모듈의 출력으로 로짓을 예측한다. 2) 핵심 메커니즘: MagicNorm는 각 모듈 내부 PreNorm 블록으로 구성되되 출구에서 최종 Norm을 적용한 합산으로 zn을 계산한다. forward에서 활성화 분산을 억제하고 backward에서 K 단계의 경사를 제한하는 TBPTT 효과를 흡수한다. 3) 학습 전략: x = (xq, xa)에서 −log P(xa|xq) 손실을 계산하고 PrefixLM 마스크를 적용해 인스트럭션 토큰은 양방향 주의를 가지며 응답 토큰은 순차적으로 생성한다. 4) 구현 세부: HAB(Hidden) 모듈은 16 layers, hidden size 1536, head size 128, context size 4096, RoPE, RMSNorm 등을 사용한다. 배치 크기 196,608 토큰, 학습률 2.2e-4, EMA 0.9999, 플랫폼은 PyTorch FSDP를 활용한다. 5) 데이터 및 샘플링: 176.5B 토큰 원천에서 40B 고정 토큰으로 stratified 샘플링을 적용하고, 4개의 주 조건 태그를 이용해 인스트럭션 스타일을 제어한다. 6) 평가 절차: 벤치마크 비교는 FLOPs를 고정한 비교와 40B 토큰 budget에서의 성능을 확인한다.

주요 벤치마크에서 HRM-Text 1B가 60.7(MMLU), 81.9(ARC-C), 63.4(Hella) 등의 점수를 기록하고, 2–7B 규모의 오픈 모델과 유사한 성능을 달성하였다. 같은 토큰 예산으로도 40B 토큰 데이터에서의 학습이 가능하며, 96–432×의 compute 절감과 100–900×의 토큰 절감을 달성하였다. ablation 연구에서 task-completion objective, PrefixLM, HRM 아키텍처의 조합이 모든 벤치마크에서 성능 향상을 주도하며, HRM으로의 전환이 최종 성능을 일관되게 상향시킨다. 효과적 깊이 분석은 HRM이 깊은 층에서도 활발한 representational 변화와 안정적 출력을 유지함을 보여주며, logit lens KL 분석에서 깊은 층의 예측이 최종 분포에 여전히 기여함을 확인한다.

HRM-Text는 2-layer 이상의 L 모듈과 H 모듈로 구성된 듀얼 타임스케일 재귀를 채택한다. 각 사이클에는 3회 L 모듈 업데이트 후 1회 H 모듈 업데이트가 수행되며, 최종 H 모듈에서 로그잣을 얻는다. MagicNorm은 N개의 재귀 forward 단계에서 각 단계마다 Norm를 적용한 후 누적 합을 더하는 방식으로 작동한다. backward pass는 TBPTT의 K-단계로 제한되며, K는 초기 2에서 시작해 5로 확장된다. Gradient는 L 내부의 PreNorm 경로를 통해 흐르고, H 블록의 정상화가 forward 경로를 안정화한다. 학습 목표는 −log P(xa|xq)이며 PrefixLM 마스크를 통해 인스트럭션(xq)은 양방향으로, 응답(xa)은 순차적으로 생성된다.