[논문 리뷰] Mechanistic Design and Scaling of Hybrid Architectures

Michael Poli, Armin W. Thomas외 10명|arXiv (Cornell University)|2024. 03. 26.

Architecture and Computational Design인용 수 4

한 줄 요약

이 논문은 기계적 아키텍처 설계(MAD)를 소개한다. MAD는 압축 및 재활용과 같은 합성 토큰 조작 작업을 활용해 하이브리드 딥 러닝 아키텍처의 확장 행동을 예측하는 빠르고 확장 가능한 파이프라인이다. 이러한 작업에 대해 소규모 모델을 훈련시킴으로써 MAD는 계산 최적화 및 과도 훈련 환경에서 최신 기술(SOTA) 모델(예: 트랜스포머, 머미바, 하이센라)을 초월하는 높은 성능을 보이는 아키텍처를 식별한다. MAD 점수와 실제 확장 성능 간에 강한 상관관계가 존재한다.

ABSTRACT

The development of deep learning architectures is a resource-demanding process, due to a vast design space, long prototyping times, and high compute costs associated with at-scale model training and evaluation. We set out to simplify this process by grounding it in an end-to-end mechanistic architecture design (MAD) pipeline, encompassing small-scale capability unit tests predictive of scaling laws. Through a suite of synthetic token manipulation tasks such as compression and recall, designed to probe capabilities, we identify and test new hybrid architectures constructed from a variety of computational primitives. We experimentally validate the resulting architectures via an extensive compute-optimal and a new state-optimal scaling law analysis, training over 500 language models between 70M to 7B parameters. Surprisingly, we find MAD synthetics to correlate with compute-optimal perplexity, enabling accurate evaluation of new architectures via isolated proxy tasks. The new architectures found via MAD, based on simple ideas such as hybridization and sparsity, outperform state-of-the-art Transformer, convolutional, and recurrent architectures (Transformer++, Hyena, Mamba) in scaling, both at compute-optimal budgets and in overtrained regimes. Overall, these results provide evidence that performance on curated synthetic tasks can be predictive of scaling laws, and that an optimal architecture should leverage specialized layers via a hybrid topology.

연구 동기 및 목표

• 매우 넓은 설계 공간과 긴 훈련 주기로 인해 딥 러닝의 아키텍처 프로토타이핑이 매우 비싸고 비효율적인 데서 비롯된 문제를 해결하기 위해.
• 전체 스케일 훈련 이전에 새로운 아키텍처를 평가하고 최적화할 수 있는 체계적이고 확장 가능한 방법을 개발하기 위해.
• 정제된 합성 작업에서의 성능이 새로운 아키텍처의 실제 세계 확장 행동을 예측할 수 있는지 조사하기 위해.
• 특정 계산 원리(예: 컨볼루션, 순환, MoE 등)를 조합한 하이브리드 아키텍처를 탐색하여 확장 효율성과 성능을 향상시키기 위해.
• 모델 능력의 기계적 이해에 기반한 계산 효율적인 아키텍처 탐색 프레임워크를 구축하기 위해.

제안 방법

• 특정 모델 능력(예: 컨텍스트 기반 학습, 시퀀스 유지)을 탐사하기 위해 압축, 재활용, 기억 등 합성 토큰 조작 작업의 세트를 설계한다.
• 반복적 및 비반복적 원리 간의 공정한 비교를 보장하기 위해 정규화된 상태 차원을 사용하여 소규모(예: 2-블록 모델)에서 아키텍처를 평가한다.
• 고립된 프록시 작업을 통해 집합적 MAD 점수를 계산하여, 다양한 단위 테스트에서 아키텍처의 능력 프로파일을 표현한다.
• 다양한 아키텍처에서 500개 이상의 언어 모델(70M–7B 파라미터)을 훈련시어 계산 최적화 및 상태 최적화 확장 법칙 분석을 수행한다.
• MAD 점수와 The Pile에서의 계산 최적화 퍼플렉서티 간의 상관관계를 분석하여 예측 능력을 검증하며, 일관된 평가를 위해 등가 FLOP 그룹을 사용한다.
• 특정 능력을 활용하는 데 특화된 계산 원리(예: 머미바 + Mixture-of-Experts)를 번갈아 배치하는 스크래치 방식을 통해 하이브리드 아키텍처를 설계한다.

실험 결과

연구 질문

• RQ1 소규모, 합성된 토큰 조작 작업은 대규모 모델 성능과 확장 법칙을 신뢰성 있게 예측할 수 있는가?
• RQ2 특정 계산 원리(예: 게이트 컨볼루션, 순환)를 조합한 하이브리드 아키텍처는 트랜스포머와 같은 균일한 아키텍처보다 확장 효율성이 뛰어나게 되는가?
• RQ3 MAD 합성 작업에서의 성능과 실제 언어 모델링에서의 계산 최적화 퍼플렉서티 간에 강한 상관관계가 존재하는가?
• RQ4 MAD 파이프라인이 특정 하위 작업에서 뛰어난 성능을 보이는 원리를 식별함으로써 반복적인 아키텍처 개선을 이끌 수 있는가?
• RQ5 상태 최적화 확장 법칙은 계산 최적화와 어떻게 다를 수 있으며, 이는 추론 효율성에 어떤 영향을 미치는가?

주요 결과

• 모든 FLOP 예산에서 집합적 MAD 점수는 계산 최적화 퍼플렉서티와 강한 선형 상관관계를 보이며, MAD가 확장 행동 예측의 유용한 대체 지표임을 검증한다.
• 특히 머미바와 MoE 또는 다른 원리를 조합한 스트립드 방식의 하이브리드 아키텍처는 트랜스포머++, 하이센라, 머미바와 같은 최신 기술 모델보다 뛰어난 확장 성능을 달성한다.
• MAD를 통해 식별된 최고 성능 아키텍처는 계산 최적화 환경과 과도 훈련 설정 모두에서 최신 기술 모델을 뛰어넘는 성능을 보였다.
• 기억 작업에서의 성능는 아키텍처 간에 거의 변하지 않았고, 확장 성능와 상관관계도 없었으며, 이 작업이 이 기준으로서의 분류 능력이 제한적임을 시사한다.
• 상태 최적화 확장 법칙 분석을 통해 고정 상태 블록(예: 머미바)이 추론 메모리 및 지연 시간 측면에서 더 높은 효율성을 달성할 수 있음을 확인하였으며, 이는 아키텍처 설계에서의 상호 보완적 트레이드오프를 반영한다.
• MAD에 기반한 반복적 개선(예: 하이센라 변형 개선)이 2×10¹⁹ FLOPs에서 The Pile에서의 퍼플렉서티 향상으로 직접적으로 이어졌으며, 이는 파이프라인의 실용적 유용성을 확인한다.