HASTE: Hardware-Aware Dynamic Sparse Training for Large Output Spaces

기술 요약: HASTE – 대규모 출력 공간을 위한 하드웨어 인지형 동적 희소 훈련(Hardware-Aware Dynamic Sparse Training)

문제 정의
극단적 다중 레이블 분류(Extreme Multi-Label Classification, XMC)는 수백만 개의 레이블이 포함된 출력 공간을 대상으로 모델을 훈련하는 과정을 포함한다. 출력층의 산술 복잡도를 줄이기 위해 희소성 기반 방법들이 제안되어 왔으나, 이들은 실제 환경에서 비례하는 속도 향상을 달amel하지 못하는 경우가 많다. 주요 병목 현상은 불규칙한 메모리 액세스 패턴, 밀집되고 연속적인 블록에 최적화된 현대 GPU의 하드웨어 활용도 저하, 그리고 비결합(uncoalesced) 읽기로 인한 "메모리 벽(memory wall)" 현상이다. 또한, 희소한 연결성을 가진 롱테일(long-tailed) 분포에서 훈련할 경우 약한 그래디언트 신호로 인해 불안정성이 발생하기 쉬우며, 이는 종종 자체적인 계산 오버헤드와 잠재적인 그래디언트 충돌을 유발하는 보조 목적 함수를 필요로 한다.

방법론
저자들은 새로운 준구조적 희소성 패턴과 데이터 기반 최적화 전략을 결합한 프레임워크인 HASTE를 제안한다.

1. 그룹 공유 고정 팬-인 희소성 (Group-Shared Fixed Fan-In Sparsity):
   표준 고정 팬-인 희소성(각 레이블이 독립적인 입력 피처 세트를 갖는 방식)이나 경직된 블록 희소성과 달리, HASTE는 "그룹 공유" 방식을 도입한다.

   • 구조: 레이블은 (훈련 인스턴스에서 유도된 레이블 임베딩을 사용하여) 의미론적 유사성에 따라 그룹으로 분할된다. 그룹 내의 모든 레이블은 동일한 희소 입력 피처 인덱스(팬-인 패턴)를 공유하지만, 독립적인 가중치 벡터는 유지한다.
   • 하드웨어 정렬: 이 설계는 "gather-once and dense MMA" 실행 패턴을 가능하게 한다. 레이블을 그룹화함으로써, 그룹에 대한 입력 피처를 한 번에 셰어드 메모리(shared memory)로 가져와 그룹 내 모든 레이벨이 재사용할 수 있도록 한다. 이는 커스텀 CUDA 커널을 통한 텐서 코어(Tensor Core) 프리미티브 사용을 용이하게 하여, 산술적 리덕션을 실제 벽시계 시간(wall-clock time)의 속도 향상으로 전환한다.
   • 메모리 효율성: 이 접근 방식은 인덱스 메모리 오버헤드를 극적으로 줄인다. 모든 레이블에 대해 인덱스를 저장하는 대신, 그룹당 한 번씩 인덱스를 저장함으로써 그룹 크기에 비례하여 오버헤드를 감소시킨다.

2. 헤드-테일 분해 (Head–Tail Decomposition, 보조 함수 없는 최적화):
   보조 목적 함수에 의존하지 않고 롱테일 분포에서의 훈련 불안정성을 해결하기 위해, HASTE는 레이블 공간을 다음과 같이 분해한다:

   • 헤드(Head): 빈도가 높은 소수의 레이블은 표준 밀집 분류기(dense classifier)를 사용하여 모델링한다.
   • 테일(Tail): 나머지 대다수의 드문 레이블은 그룹 공유 희소 레이어를 통해 처리된다.
   • 메커니즘: 두 브랜치는 동일한 인코더를 공유하지만 별도의 경량 피처 프로젝션을 사용한다. 밀집 헤드는 훈련 과정 중 인코더에 일관되고 강력한 그래디언트 신호를 제공하여, 별도의 보조 감독 브랜치 없이도 희소 테일의 학습 과정을 안정화한다.

주요 기여

• 새로운 희소성 패턴: 하드웨어 효율적인 구조적 공유 액세스와 독립적인 피처 선택의 표현력 사이의 균형을 맞추는 그룹 공유 고정 팬-인 희소성을 도입하였다.
• 하드웨어 인지형 커널: 모여진 피처 타일(feature tiles)에 대해 밀집 행렬 곱셈(GEMM)을 수행하여 텐서 코어의 효율적인 사용을 지원하는 커스텀 CUDA 커널을 개발하였다.
• 보조 함수 없는 안정성: 레이블의 내재된 빈도 분포를 활용하여 훈련을 안정화하는 분해 전략을 통해, 보조 목적 함수 및 그에 따른 오버헤드 없이 롱테일 환경에서의 학습을 안정화하였다.
• 실증적 검증: 최대 860만 개의 레이블에 달하는 대규모 벤치마크를 통해, HASTE가 기존의 희소 베이스라인들과 비교하여 속도 및 메모리 효율성 면에서 크게 앞서면서도 밀집 모델과의 성능 격차를 좁힌다는 것을 입증하였다.

실험 결과

• 속도 향상: NVIDIA A100 GPU에서의 마이크로 벤치마크 결과, HASTE는 표준 고정 팬-인 희소성(SPARTEX) 대비 순방향 패스(forward pass)에서 최대 4.4배, 역방향 패스(backward pass)에서 최대 25배의 속도 향상을 달성했다. 이는 FLOPs가 일치하는 밀집 병목 모델의 성능에 근접하게 작동한다.
• 정확도: Amazon-670K, Amazon-3M, LF-Paper2Keywords-8.6M을 포함한 벤치마크에서, HASTE는 최신 희소 베이스라인(예: SPARTEX, BLOCK SPARSE)과 비교하여 Precision@k (P@k) 및 Propensity-Scored Precision@k (PSP@k) 측면에서 대등하거나 더 나은 성능을 보인다.
• 메모리 효율성: 이 방법은 피크 훈련 메모리를 크게 줄이며(예: AmazonTitles-670K에서 SPARTEX의 5.0 GiB 대비 3.2 GiB), 에포크(epoch) 시간을 단축한다.
• 그룹 크기 트레이드오프: 실험에 따르면, 더 나은 표현력의 유연성을 위해 작은 그룹 크기(예: 16)를 사용하는 것이 정확도 면에서 약간 더 유리한 반면, 더 큰 그룹(예: 64)은 피처 재사용을 높여 커널 효율성을 개선한다는 트레이드오프 관계가 존재한다.

의의 및 주장
본 논문은 HASTE가 희소성의 이론적 산술 이점과 현대 GPU 하드웨어의 실제적 현실 사이의 간극을 성공적으로 메웠다고 주장한다. 그룹 공유를 통해 희소성 패턴에 구조를 도입하고 이를 XMC 데이터의 롱테일 특성에 정렬함으로써, FLOPs 감소를 실질적인 벽시계 시간 이득으로 전환하였다. 저자들은 이 방식이 훈련 시간과 메모리 점유율을 모두 줄이면서도 경쟁력 있는 정확도를 유지함으로써, 고정된 컴퓨팅 예산 하에서 수백만 개의 레이블을 가진 모델의 훈련을 더욱 용이하게 만든다고 주장한다. 또한, 효율성은 개선되었으나 이것이 롱테일 데이터셋의 데이터 편향이나 불완전한 레이블링 문제를 본질적으로 해결하는 것은 아니므로, 사용자는 테일 레이블의 성능에 대해 주의를 기울여야 한다고 명시하였다.