선형 회귀 인컨텍스트 학습에서 선형·이차 어텐션의 성능 비교

초록

본 연구는 선형 회귀 ICL(인컨텍스트 학습) 벤치마크에서 이차(소프트맥스) 어텐션과 선형(커널) 어텐션 트랜스포머의 학습 품질, 수렴 속도, 일반화 능력을 체계적으로 비교한다. 깊이(1·3·6층)와 학습 단계별 성능을 분석한 결과, 이차 어텐션이 최종 MSE와 분산 이동에 있어 우수하지만, 선형 어텐션은 더 빠른 수렴과 비슷한 깊이 확장 효과를 보인다. 또한, 두 모델 모두 깊이가 증가할수록 ICL 능력이 크게 향상되며, 이차 어텐션이 비등방성 데이터에 대해 약간 더 강인함을 보인다.

상세 분석

이 논문은 Garg et al.이 제시한 선형 회귀 ICL 태스크를 재현하고, 두 종류의 어텐션 메커니즘을 동일한 하이퍼파라미터 설정 하에 비교한다. 입력 차원 dₓ=5, 임베딩 차원 d_model=256, 헤드 수 4, MLP 확장 비율 4 등으로 모델 규모를 통일했으며, 1·3·6층 깊이에서 각각 학습시켰다. 이차 어텐션은 기존 GPT‑2 스타일의 소프트맥스 기반으로 O(T²) 복잡도를 갖고, 학습률 1e‑4, 배치 32, 30 k 스텝을 사용했다. 선형 어텐션은 Katharopoulos et al.의 커널화 기법을 채택해 O(T d²) 복잡도로 구현했으며, ReLU² 특성 맵을 적용해 비음수성을 확보하고, 학습률 3e‑4, 배치 64, 층당 7.5k~10k 스텝으로 훈련했다.

성능 측면에서, 6층 이차 어텐션은 정규화 MSE 0.0203±0.0016으로 거의 최적에 근접했으며, 3층에서는 0.0502±0.0063, 1층에서는 0.7976±0.0543으로 급격히 개선된다. 반면 선형 어텐션은 6층에서 0.0102±0.0012, 3층에서 0.0366±0.0021, 1층에서 0.5946±0.0483을 기록했는데, 특히 6층 모델이 최종 MSE에서 이차 어텐션보다 약간 낮은 값을 보였지만, 전체적인 추세는 이차 어텐션이 더 안정적이고 낮은 오차를 유지한다는 것이다.

수렴 속도 분석에서는 90% 최종 MSE 도달에 필요한 샘플 수를 비교했을 때, 선형 어텐션이 1·3층에서 256 k 샘플, 6층에서 480 k 샘플로 이차 어텐션(각각 224 k, 688 k, 800 k)보다 약 1.5~2배 빠르게 수렴한다. 이는 선형 어텐션이 제한된 가설 공간과 낮은 파라미터 수 덕분에 초기 학습 단계에서 급격히 손실을 감소시키는 특성으로 해석된다.

일반화 측면에서는 비등방성 공분산 Σ를 가진 데이터셋에 대해 두 모델 모두 성능 저하가 있었지만, 이차 어텐션이 상대적으로 작은 손실 증가(예: 6층에서 0.0365→0.0398) 를 보이며, 선형 어텐션은 0.0302→0.0328 로 비슷한 비율이지만 절대값이 약간 더 높다. 이는 이차 어텐션이 더 풍부한 표현력을 통해 입력 분포 변화에 대한 적응력이 약간 우수함을 시사한다.

깊이 확장 효과는 양쪽 모두에서 뚜렷하게 나타난다. 1→3층 전이에서 이차 어텐션은 89% 손실 감소, 선형 어텐션은 87% 감소를 기록한다. 3→6층에서는 이차 어텐션이 추가 59% 감소, 선형 어텐션은 68% 감소를 보이며, 선형 모델이 깊이 증가에 따라 더 큰 상대적 이득을 얻는 경향이 있다. 다만 6층 선형 모델은 3층 모델보다 학습 스텝이 더 많이 필요해 최적점에 도달하는 시점이 늦어지는 현상이 관찰되었다.

전체적으로, 이 논문은 선형 어텐션이 계산 효율성과 빠른 수렴 측면에서 실용적 장점을 제공하지만, 최종 정확도와 분산 이동에 대한 강인성에서는 이차 어텐션이 여전히 우위에 있음을 보여준다. 또한, 모델 깊이가 ICL 능력에 결정적인 영향을 미치며, 두 어텐션 방식 모두 깊이 증가에 따라 급격히 성능이 향상되는 점을 확인했다. 이러한 결과는 대규모 언어 모델에서 메모리·시간 복잡도를 절감하기 위해 선형 어텐션을 도입할 경우, 충분한 깊이와 적절한 커널 선택이 필요함을 시사한다.