본 논문은 Transformer 모델의 핵심인 멀티-헤드 어텐션 메커니즘을 효율성을 향상시키면서 기존의 정확도 수준을 유지하거나 능가하도록 업그레이드하는 것을 목표로 한다. 멀티-헤드 어텐션을 합의 형태로 표현할 수 있다는 점에 착안하여, 모든 어텐션 헤드가 동일한 중요성을 갖는 것은 아니라는 점을 고려하여 Mixture-of-Head (MoH) 어텐션이라는 새로운 아키텍처를 제안한다. MoH는 어텐션 헤드를 Mixture-of-Experts (MoE) 메커니즘의 전문가로 취급한다. MoH는 각 토큰이 적절한 어텐션 헤드를 선택할 수 있게 함으로써 정확도를 저하시키거나 파라미터 수를 늘리지 않고 추론 효율성을 높이는 장점을 갖는다. 또한, 표준 합산 대신 가중치 합산을 사용하여 어텐션 메커니즘에 유연성을 더하고 성능 향상의 가능성을 열어준다. ViT, DiT, 그리고 LLMs에 대한 광범위한 실험을 통해 MoH가 어텐션 헤드의 50%-90%만 사용하면서도 멀티-헤드 어텐션보다 우수한 성능을 보임을 보여준다. 특히, LLaMA3-8B와 같은 사전 훈련된 멀티-헤드 어텐션 모델을 MoH 모델로 추가적으로 계속해서 미세 조정할 수 있음을 보여주며, MoH-LLaMA3-8B는 14개의 벤치마크에서 평균 64.0%의 정확도를 달성하여 어텐션 헤드의 75%만 사용하면서 LLaMA3-8B보다 2.4% 향상된 성능을 보였다.
