Prism Transformer는 레이어 간 어텐션 헤드 수를 가변적으로 조절하는 점진적 헤드 스케줄을 도입합니다. 이는 초기 레이어에서 적고 넓은 헤드로 시작하여 깊이에 따라 단조 증가시킵니다. 이 접근 방식은 아키텍처 오버헤드를 추가하지 않고 초기와 후기 레이어의 구조적 필요를 해결함으로써 표준 균일 할당에 도전합니다.

  • 초기 레이어는 풍부한 국소 패턴 포착을 위해 넓은 헤드 부분공간(dh=256)을 활용하고, 후기 레이어는 특수화된 분해를 위해 많은 좁은 헤드를 사용합니다.
  • 가중치 행렬은 표준 dmodel×dmodel 형태를 유지하여 파라미터 수를 중립적으로 유지합니다.
  • 총 FLOPs는 헤드 수에 대해 수학적으로 불변이므로 계산 중립성을 보장합니다.
  • 2의 거듭제곱인 헤드 차원(dh ∈ {256, 128})은 처리량 중립성을 위해 Tensor Core 정렬을 유지합니다.
  • 결과는 동일한 토큰/초와 실행 시간을 가진 균일 베이스라인과 비교하여 모든 규모에서 검증 손실이 낮음을 보여줍니다.
  • 이 모델은 PIQA, WinoGrande, HellaSwag, ARC-Easy를 포함한 벤치마크에서 향상 또는 동등한 성능을 달성합니다.

레이어별 어텐션 거리 분석은 이 이득이 구조적임을 확인합니다. 초기 Prism 레이어는 전역 통합으로 전환하기 전에 더 국소적으로 어텐션을 수행합니다. 구현에는 헤드 수를 레이어 종속적으로 만들기 위해 각 어텐션 레이어당 1줄의 변경만 필요합니다.