최근 공개된 연구 논문에서 '프리즘 트랜스포머(Prism Transformer)'라는 새로운 아키텍처가 기존 트랜스포머 모델의 고질적인 문제점을 해결하며 성능 향상을 이끌어냈습니다. 기존 트랜스포머는 멀티 헤드 어텐션 메커니즘에서 모든 레이어에 걸쳐 히든 차원(hidden dimension)을 동일하게 분할하여 어텐션 헤드에 할당했습니다. 연구진은 이러한 균일한 할당 방식이 구조적 병목 현상을 일으켜, 초기 레이어의 헤드들이 복잡하고 고차원적인 문맥 패턴을 효과적으로 포착하지 못한다고 지적했습니다.
프리즘 트랜스포머는 이 문제를 해결하기 위해 '점진적인 헤드 스케줄(progressive head schedule)' 방식을 도입했습니다. 이는 레이어가 깊어질수록 어텐션 헤드의 개수를 단조적으로 증가시키는 방식입니다. 즉, 초기 레이어에서는 적은 수의 매우 넓은(wide) 헤드를 사용하여 복잡하고 지역적인(local) 구성 패턴을 포착하고, 후기 레이어에서는 더 많은 수의 좁은(narrow) 헤드를 배치하여 이러한 패턴들을 전문화된 언어적 특징으로 분해합니다. 이러한 구조적 변화는 놀랍게도 파라미터(parameter)나 연산량(FLOPs)의 증가 없이 이루어지며, 학습 및 추론(inference) 오버헤드(overhead)도 발생하지 않습니다.
이러한 비균일한 서브스페이스(subspace) 할당 방식은 표준 트랜스포머의 잠재된 역량을 끌어내 모델 용량을 더욱 효과적으로 활용하게 합니다. 1.24억, 3.54억, 7.57억 개의 파라미터를 가진 세 가지 모델 규모에서 프리즘 트랜스포머는 일관되게 기존 균일 방식의 베이스라인 모델보다 우수한 성능을 보였습니다. 이는 검증 손실(validation loss) 감소와 함께 PIQA, HellaSwag, ARC-Easy, WinoGrande 등 다양한 제로샷(zero-shot) 벤치마크에서 지속적인 성능 향상으로 이어졌습니다. 이 연구는 기존 트랜스포머 아키텍처의 효율성을 극대화할 수 있는 새로운 방향을 제시하며, 향후 대규모 언어모델(LLM) 개발에 중요한 시사점을 제공합니다.