AI Researcher & Engineer

5 minute read

https://arxiv.org/abs/2310.07707

MatFormer: Nested Transformer for Elastic Inference

구글에서 소규모 모델에서 사용했다고 말하는 matformer를 리뷰할것이다.

초록
파운데이션 모델(Foundation Models)은 대규모 멀티 가속기 클러스터부터 리소스가 제한된 독립형 모바일 기기까지 다양한 추론 제약 조건을 가진 환경에서 활용되고 있습니다. 하지만 이러한 모델을 훈련하는 데 드는 막대한 비용은 제공할 수 있는 모델 크기의 다양성을 제한하는 요인이 됩니다. 그 결과, 실사용자는 자신이 원하는 지연 시간(latency)이나 비용 요구 사항에 최적화되지 않은 모델을 선택할 수밖에 없는 상황에 놓이게 됩니다.

이에 우리는 MatFormer를 소개합니다.
※ MatFormer는 Matryoshka Transformer의 줄임말로, 모델이 본질적으로 중첩된(nested) 구조를 가진다는 점에서 이름이 유래되었습니다.

MatFormer는 다양한 배포 제약 조건을 고려해 탄력적인 추론(elastic inference)을 제공할 수 있도록 설계된 새로운 형태의 Transformer 아키텍처입니다. 이 모델은 표준 Transformer 구조 내에 중첩된 Feed Forward Network (FFN) 블록 구조를 통합함으로써 이를 구현합니다. 학습 시, 크기가 서로 다른 여러 중첩 FFN 블록의 파라미터를 함께 최적화함으로써 추가 연산 비용 없이 수백 개의 정확한 소형 모델들을 추출할 수 있습니다.

우리는 MatFormer의 효용성을 디코더와 인코더 등 다양한 모델 유형뿐 아니라, 언어와 비전 등 다양한 모달리티에 걸쳐 실험적으로 검증했습니다. 그 결과, 실제 환경에서도 효과적으로 활용될 수 있는 가능성을 보여줍니다.

예를 들어, 파라미터 수 8억 5천만 개(850M)의 디코더 전용 MatFormer 언어 모델(MatLM)에서 5억 8천 2백만 개(582M)에서 8억 5천만 개 사이의 여러 소형 모델을 추출할 수 있으며, 각각은 독립적으로 학습한 동등한 크기의 모델보다 더 나은 검증 손실 및 단발(one-shot) 다운스트림 성능을 보였습니다.

또한, 범용 MatFormer 기반의 비전 Transformer(MatViT) 인코더에서 추출한 소형 인코더들도 대규모 적응형 검색(adaptive large-scale retrieval)에 필요한 거리 기반 임베딩 구조(metric-space structure)를 잘 보존하는 것을 확인했습니다.

마지막으로, MatFormer에서 추출한 정확하고 일관된 서브모델을 활용한 사전 추론(speculative decoding) 기법을 통해 추론 지연 시간(inference latency)을 크게 줄일 수 있음을 입증했습니다.

자세한 내용은 프로젝트 웹사이트에서 확인할 수 있습니다.

1 서론

Figure 1: MatFormer는 Transformer의 FFN 블록에 중첩(nested) 구조를 도입하고, 모든 서브모델을 함께 학습함으로써, 수백 개의 정확한 서브모델을 자유롭게 추출할 수 있게 하여 탄력적인 추론(elastic inference)을 가능하게 한다.

대규모 파운데이션 모델(FM, Foundation Model) [49, 45, 17]은 모바일 환경의 실시간 응답부터 웹 규모 배치 처리(batch serving)를 위한 멀티 클러스터 GPU에 이르기까지, 다양한 연산 및 정확도 요구사항을 가진 환경에서 활용됩니다. 그러나 일반적인 모델 계열(family)은 서로 다른 크기의 모델을 몇 개만 독립적으로 훈련하여 제공합니다. 예를 들어, Llama-2 모델 계열은 70억(7B), 130억(13B), 340억(34B), 700억(70B) 파라미터 모델만을 제공합니다 [59]. 이로 인해 실제 사용자는 자신의 지연 시간이나 비용 제약에 비해 더 작은(그리고 일반적으로 덜 정확한) 모델을 선택해야 하는 경우가 많습니다. 또는 압축(compression)이나 프루닝(pruning)을 활용해 더 큰 모델을 주어진 예산 안에 맞춰 넣는 방법도 있지만 [19, 36, 53], 이 경우 추가적인 학습이 필요합니다.

우리는 이러한 문제를 해결하기 위해 MatFormer를 소개합니다. 이는 기본적으로 탄력성(elasticity)을 내장한 Transformer 아키텍처 [61]로, 하나의 범용(universal) 모델을 학습한 뒤, 추가 학습 비용 없이 수백 개의 소형 서브모델을 추출할 수 있도록 설계되어 있습니다 (그림 1 참조). MatFormer는 인코더와 디코더 모두에 적용 가능하며, 도메인에 무관하게 사용 가능하고, 기존 파운데이션 모델 훈련 파이프라인과도 호환됩니다.

MatFormer는 Matryoshka 표현 학습(Metryoshka Representation Learning) [34] 원리를 기반으로, 표준 Transformer 블록 내부에 중첩된 하위 구조(nested substructure)를 도입합니다. 형식적으로는 Transformer 블록들을
T₁ ⊂ T₂ ⊂ ⋯ ⊂ T_g
와 같이 정의하며, 여기서 g는 중첩 블록의 수를 의미하고, T_i ⊂ T_i+1 관계는 T_i의 파라미터가 T_i+1에 포함된다는 것을 뜻합니다. MatFormer는 Transformer의 attention 블록과 Feed Forward Network (FFN) 블록 양쪽 모두에 이러한 중첩 구조를 도입할 수 있습니다 (그림 1 참조).

예를 들어, FFN 블록의 hidden layer에 d_ff개의 뉴런이 있다고 할 때, MatFormer는 이들 뉴런에 대해 matryoshka 구조를 유도하여, T_i는 상위 m_i개의 뉴런만 포함합니다. 이때,
1 ≤ m₁ < m₂ < ⋯ < m_g = d_ff
로 설정되며, 각 단계별 서브모델에 포함되는 뉴런 수를 의미합니다. 직관적으로, 처음 m₁개의 뉴런이 가장 중요하며, 이후의 뉴런은 점차 더 큰 모델에만 포함됩니다.

기존 연구들과 달리 (2절 참고), 우리는 단지 g개의 크기(세분화 수준, granularity)에 대해 학습했음에도 불구하고, 학습 이후 지수적으로 많은 수의 서브모델을 추출할 수 있습니다. 학습된 MatFormer 블록 T₁, …, T_g를 각 층에서 조합함으로써 새로운 모델을 구성할 수 있는데, 이를 우리는 Mix’n’Match 방식이라고 부릅니다 (3.3절 참조). 예컨대, 첫 번째 층에서 T_g, 두 번째 층에서 T₂, 세 번째 층에서 T₄ 등의 조합을 선택하여 총 g_l개의 조합이 가능하며, 여기서 l은 층의 수입니다. 놀랍게도, 다양한 모델 크기와 설정에서 이러한 방식으로 추출된 모델들이 높은 정확도를 보이며, 추출된 모델의 크기에 따라 정확도도 자연스럽게 향상됨을 확인했습니다.

우리는 최대 8억 5천만 개의 파라미터를 가진 디코더 전용 언어 모델(MatLM)을 MatFormer 기반으로 학습하여 다음과 같은 결과를 관찰했습니다:

  1. g개의 세분화 수준으로 학습된 MatLM은, 각각 독립적으로 학습된 동일 크기의 모델들보다 검증 손실과 다운스트림 one-shot 성능에서 더 우수했습니다.
  2. Mix’n’Match 방식으로 추출된 모델들은, g개의 명시적 학습 모델들이 형성한 정확도-파라미터 트레이드오프 곡선 상에 존재합니다.
  3. 다양한 크기의 MatFormer 모델에 대해 손실-연산량 법칙(loss vs. compute law)이 일반 Transformer와 유사하게 유지됨을 확인했습니다.
  4. MatLM에서 추출된 서브모델은 일관성이 높아, 추론 최적화 및 다양한 규모의 환경에 배포하기에 매우 적합합니다.

우리는 ViT 기반의 MatFormer 모델(MatViT)도 추가로 실험했습니다. 예를 들어, MatViT-L/16은 ImageNet-1K에서 표준 ViT-L/16보다 더 높은 정확도를 달성했으며, 추출된 모든 서브모델 또한 각각의 독립 학습 모델과 같거나 더 나은 성능을 보였습니다. 특히, MatViT는 일관성이 매우 높기 때문에, 탄력적 이미지 검색(elastic image retrieval)을 위한 “탄력적 인코더(elastic encoder)”로 활용할 수 있습니다. 즉, 가장 큰 MatViT 모델로 인코딩된 이미지의 임베딩 공간(metric space)은, 서브모델에 의해서도 잘 보존됩니다. 이에 따라, 질의 복잡도나 시스템 부하 등의 조건에 따라, 서브모델 중 하나를 선택하여 고정된 코퍼스에서 검색을 수행할 수 있고, 이 경우 계산량을 40% 이상 줄이면서도 정확도 손실은 0.5% 이하로 유지됩니다.

주요 기여 내용은 다음과 같습니다:

  1. 표준 Transformer 내에 중첩 구조를 도입하고, g개의 세분화 수준을 최적화하여 하나의 범용 탄력 모델을 구성하는 MatFormer를 제안합니다.
  2. 추가 계산 비용 없이 파라미터 제약 내에서 최적의 서브모델을 찾아주는 간단한 휴리스틱인 Mix’n’Match 기법을 소개합니다. 이 기법은 복잡한 NAS 기법보다 우수한 결과를 보이며, 수백 개의 정확하고 일관된 서브모델을 추가 학습 없이 생성할 수 있습니다 (3절 참조).
  3. MatFormer는 디코더 전용 언어 모델(MatLM)과 비전 인코더(MatViT) 모두에 잘 일반화되며, 표준 Transformer와 유사한 수준의 성능과 확장성을 가지면서도, 훨씬 더 빠른 오토레그레시브 생성과 대규모 적응형 밀집 검색을 가능하게 합니다 (4절 참조).

? 구조적으로는 좋아 보이지만 실제로는 다음이 걱정된다는 관점:

1. 학습 안정성 문제

  • 중첩 구조(nested FFN) 안에 여러 개의 서브모델을 한꺼번에 학습하므로,
  • 특정 서브모델의 파라미터가 전체 모델에 중복 포함되면서 gradient conflict가 발생할 수 있음.
  • 특히 작은 서브모델이 성능을 내기 위해 집중된 학습을 받아야 하는데, 큰 모델이 그 파라미터를 공유하면 균형 잡힌 학습이 어렵고 오버핏/언더핏 혼란 가능성 있음.
  • 논문에서는 이를 다층 Mix’n’Match로 커버한다고 했지만, 그게 정말 안정적으로 작동할지는 별도 검증이 필요함.

2. 학습 효율성

  • **“추가 연산 없이 여러 모델을 한 번에 학습한다”**는 주장을 하는데, 사실상 g개의 granularity에 대해 loss를 따로 최적화해야 하므로,
  • 일반 Transformer 대비 학습 속도 자체는 느릴 수 있고, 메모리 사용량도 늘어날 가능성 있음.
  • 논문에서는 scaling law 측면에서 vanilla transformer와 비슷하다고 주장하지만, 학습 efficiency에 대한 absolute comparison은 부족함.

3. 모델 압축(MOE/MOA) 방식 대비 실용성

  • 말씀하신 것처럼, 큰 모델 하나만 잘 학습시키고 그걸 부분적으로 선택적으로 사용하는 Mixture of Experts (MoE), Mixture of Activations (MoA), 혹은 pruning/trimming 기법이 오히려:
    • 더 단순하고,
    • 서브모델 선택이 명시적이며,
    • 하위 모델의 품질이 예측 가능한 경우가 많음.
  • 반면 MatFormer는 “서브모델이 잘 작동한다”고 주장하지만, 그 신뢰성과 일관성은 실제 환경에서는 검증이 더 필요함.

Comments

You May Also Enjoy