MaskLLM

MaskLLM은 각 가중치를 손으로 평가하는 대신 마스크 자체를 학습한다. 유효한 N:M 패턴 집합을 분포로 다루고 학습 중에 표본을 뽑으므로, 모델은 데이터셋에 가장 알맞은 마스크로 수렴한다. 이 학습 단계가 마스크를 한 번 고르고 끝내는 일회성 방식과 근본적으로 다른 점이다.

이 기법은 이산 샘플링 기술을 빌려 마스크 선택을 엔드투엔드로 학습 가능하게 만든다. Gumbel Softmax는 본래 흑백으로 갈리는 딱딱한 선택에 그래디언트를 통과시킨다. 그래서 N:M 마스크 위의 분포를 추정이 아니라 실제 데이터로 최적화할 수 있다.

프로젝트에는 바로 쓸 수 있는 마스크가 들어 있어 학습 없이도 시험해 볼 수 있다. LLaMA-2 7B, LLaMA-3 8B, LLaMA-3.1 8B용 마스크가 공개됐고, 각각 C4의 일부로 학습됐다. 스크립트가 Hugging Face에서 내려받아 평가용으로 공식 모델에 적용한다.

MaskLLM이 겨냥하는 것은 N:M 희소성, 대개 연속한 가중치 네 개에 0이 아닌 값 두 개를 남기는 2:4 패턴이다. 이 형식이 중요한 이유는 하드웨어 가속에 맞아떨어지기 때문이다. 작아진 모델은 자리만 덜 차지하는 게 아니라, 알맞은 GPU에서는 더 빠르게 돌아간다.

마스크를 배우는 같은 확률 모델이 마스크를 한 영역에서 다른 영역으로 옮긴다. 마스크 분포를 한 번 학습하면 새 작업이나 데이터셋으로 옮길 때 드는 품이 처음부터 시작하는 것보다 적다. 저자들은 이를 전이성이라 부른다. 모델을 여러 상황에서 예리하게 유지해야 할 때 도움이 된다.

본격적인 학습을 위해 LLaMA 모델을 여러 GPU에 나누는 Megatron-LM 워크플로가 함께 제공된다. 리포지터리에는 체크포인트 변환, 데이터 토큰화, SparseGPT나 Wanda, 크기 기반 프루닝으로 사전 마스크 생성, 마스크 학습, 공개 전 체크포인트 깎기까지 처리하는 스크립트가 있다.