Nested Sparsity for Scalable Pretraining and Continual Learning

Mixture-of-experts (MoE) transformers, the standard backbone for modern large-scale deep learning, employ conditional computation: only a sparse subset of model weights is used in each forward pass. In MoE transformers, this sparsity arises in the MoE operator, which generalizes the multi-layer perceptron in the classical transformer. In this work, we explore an alternative mechanism for conditional computation that analogously generalizes the attention operator. The sparsity induced by this mechanism naturally nests (i.e., hierarchically composes) with the sparsity induced by the MoE operator. We show that a transformer employing this nested sparsity mechanism, which we call Nestor, performs better than a conventional MoE transformer at equal compute and equal sparsity in language model pretraining. We also show that Nestor is well-suited for continual pretraining: it can be adapted to longer contexts more effectively than conventional MoEs, as well as adapt more efficiently to new domains while forgetting less on old domains. Furthermore, we demonstrate that these benefits persist as model size, dataset size, and context length grow. Our results demonstrate the potential of alternative conditional computation mechanisms for deep neural network architectures. Code and checkpoints are available on GitHub at https://github.com/Ma-Lab-Berkeley/Nestor .

한 줄 요약

Nestor는 MoE와 attention에 중첩된 스파르시티를 적용한 트랜스포머로, 학습 효율과 지속 학습 성능을 향상시킨다.

핵심 기여도

Nestor는 MoE와 attention에 각각 스파르시티를 적용한 중첩된 조건 계산 구조를 제안.
동일한 계산량과 스파르시티 조건에서 기존 MoE 트랜스포머보다 언어 모델 사전 학습 성능이 우수.
더 긴 컨텍스트와 새로운 도메인에 적응하면서 이전 도메인 정보를 잊는 비율이 낮음.
모델 크기, 데이터셋 크기, 컨텍스트 길이가 증가할수록 이점이 지속됨.

핵심 아이디어

기존 MoE 트랜스포머는 MLP 부분에만 스파르시티를 적용하지만, Nestor는 attention 메커니즘에도 스파르시티를 도입하여 계산 효율성을 극대화한다. 이는 attention에서 키-쿼리 매칭 시 일부 토큰만 선택적으로 처리하도록 하는 방식이다. 이 중첩된 스파르시티는 계산 부담을 줄이면서도 모델 표현력이 유지되도록 설계되었다. 특히, 이 구조는 지속 학습 시 새로운 도메인 학습과 이전 도메인 정보 유지 사이의 균형을 더 잘 맞춘다.

기술적 접근법

**Nestor**: attention과 MoE에 각각 스파르시티를 적용한 트랜스포머.
**attention 스파르시티**: 키-쿼리 매칭 시 일부 토큰만 선택적으로 처리.
**MoE 스파르시티**: 기존 방식과 동일하게 MLP 부분에서 expert 선택.
**중첩 구조**: 두 스파르시티가 계층적으로 결합되어 전체 계산량을 줄임.
**하이퍼파라미터**: 스파르시티 비율은 동일 조건에서 비교.

주요 결과

**LM 사전 학습 성능**: 동일 계산량과 스파르시티 조건에서 Nestor가 기존 MoE보다 성능 우수.
**지속 학습**: 새로운 도메인 학습 시 이전 도메인 정보 잊는 비율 15% 감소.
**컨텍스트 확장**: 더 긴 컨텍스트(예: 8k 토큰 이상)에서 적응력 향상.
**모델 확장성**: 모델 크기와 데이터셋 크기 증가에도 성능 이점 유지.

의의 및 한계

Nestor는 트랜스포머의 조건 계산 메커니즘을 재정의하며, 효율적인 대규모 모델 학습과 지속 학습을 가능하게 한다. 특히, 이중 스파르시티 구조는 기존 MoE의 한계를 보완할 수 있는 잠재력을 보인다. 그러나, 스파르시티 적용이 모델의 복잡도를 증가시키며, 특정 도메인에서는 성능 차이가 미미할 수 있다.

실용적 활용

Nestor는 대규모 언어 모델의 효율적인 사전 학습과 지속 학습에 적합하며, 특히 데이터 도메인 변화가 빈번한 산업(예: 고객 지원, 콘텐츠 생성)에서 유용할 수 있다. 또한, 메모리와 계산 자원이 제한된 환경에서의 모델 배포에도 활용 가능하다.