Transformer의 주·키 정규화에 Lp 노름을 도입한 향상된 QKNorm

초록

본 논문은 기존 QKNorm에서 사용하던 L2 정규화를 일반화하여 Lp 노름( p≥1) 기반 정규화 기법을 제안한다. nanoGPT 기반 작은 문자 수준 언어 모델에 p∈{1.0,1.5,…,4.0}를 적용해 교차 검증 실험을 수행했으며, p>2인 경우 검증 손실이 더 낮고 수렴 속도가 빨라짐을 확인했다. 계산 비용은 거의 변하지 않아 실용성이 높다.

상세 분석

본 연구는 Transformer의 핵심 연산인 스케일드 닷-프로덕트 어텐션에서 발생하는 로그스케일 폭발 문제를 완화하기 위해 Query‑Key 정규화(QKNorm)를 확장한다. 기존 QKNorm은 ℓ2 정규화와 학습 가능한 스칼라 α만을 사용했으며, 이는 모든 쿼리·키 벡터를 단위 구면에 투사한다. 저자들은 ℓp 정규화( p≥1)를 도입함으로써 정규화된 벡터의 기하학적 형태를 p값에 따라 조절할 수 있게 하였다. p가 클수록 절대값이 큰 차원에 가중치가 집중되어 “스파이크” 형태의 어텐션 분포를 유도하고, p→∞에서는 최대 절대값 차원 하나만 고려하게 된다. 이러한 특성은 어텐션이 집중해야 할 특징 차원을 선택적으로 강조하거나 억제하는 새로운 하이퍼파라미터 역할을 한다.

실험에서는 nanoGPT‑style 디코더만을 사용해 Tiny Shakespeare 데이터셋을 문자 단위로 학습시켰다. 10‑fold 교차 검증과 p값 7가지 조합으로 총 70번의 학습을 수행했으며, 검증 교차 엔트로피 손실을 기준으로 성능을 평가했다. 결과는 p=2.5, 3.0, 3.5, 4.0에서 p=2(기본 QKNorm)보다 최소 손실이 각각 1.373, 1.365, 1.362, 1.357로 유의하게 낮았다. 또한 손실 최소점에 도달하는 이터레이션 수가 감소해 수렴 속도가 빨라졌다. 훈련 시간은 p값에 관계없이 평균 360~363초로 차이가 거의 없었으며, 이는 ℓp 정규화 연산이 기존 ℓ2 정규화와 비교해 계산 복잡도에서 실질적인 오버헤드를 발생시키지 않음을 의미한다.

논문의 토론에서는 ℓp 정규화가 어텐션의 “특징 스팬”을 제어한다는 해석을 제시한다. p가 증가하면 높은 절대값을 가진 차원에만 주목하게 되어 모델이 더 적은 수의 특징에 집중한다는 가설을 세웠으며, 실험 결과는 이러한 가설을 뒷받침한다. 그러나 실험이 매우 제한된 설정(작은 문자 수준 모델, 단일 GPU, 단일 데이터셋)에서만 수행되었으며, 대규모 언어 모델이나 다른 모달리티에 대한 일반화 가능성은 검증되지 않았다. 또한 p값이 너무 크게 설정될 경우(예: p→∞) 정보 손실 위험이 존재함에도 불구하고 해당 극한 상황에 대한 분석이 부족하다.

향후 연구에서는 (1) 다양한 데이터셋·모델 규모에 대한 확장 실험, (2) ℓp 정규화와 기존 LayerNorm·Peri‑LN 등 다른 정규화 기법의 조합 효과, (3) p값을 학습 가능한 파라미터로 두어 자동 튜닝하는 메커니즘, (4) ℓp 정규화가 어텐션 헤드 간 상관관계와 표현 다양성에 미치는 영향을 이론적으로 분석하는 방향을 제안한다.