프리트레인 트랜스포머의 테스트 시점 적응성과 분포 강건성

초록

본 논문은 다양한 난이도의 과제들로 구성된 혼합 분포에서 사전 학습된 트랜스포머가, 테스트 시점에 고정된 난이도 β를 갖는 과제들에 대해 χ² 제한 안에서 발생하는 분포 이동에 강인하면서도 난이도에 따라 수렴 속도를 자동 조절한다는 이론적 결과를 제시한다. 비모수 회귀와 다중 인덱스 모델을 대상으로, 충분히 큰 트랜스포머와 충분한 사전 학습 데이터가 주어지면 최적의 수렴률을 달성하고, 테스트 분포에 대한 사전 지식이 있더라도 이를 초월하는 성능을 보임을 증명한다.

상세 분석

논문은 먼저 사전 학습 prior π를 α∈𝒜에 대한 가중합 형태로 정의하고, 각 π_α가 특정 난이도 α를 갖는 과제 분포라고 가정한다. 테스트 시점에는 고정 난이도 β의 과제들만을 포함하는 분포 μ가 존재하며, μ와 π_β 사이의 χ² 발산이 κ 이하라는 제약을 둔다. 이 설정은 실제 LLM이 다양한 난이도의 프롬프트를 사전 학습하고, 특정 도메인에 집중된 테스트 환경에 노출될 때 발생하는 분포 이동을 수학적으로 모델링한다는 점에서 의미가 크다.

핵심 이론적 기여는 Proposition 1에서 제시된 ICL excess risk의 상한식이다. 위험을 두 부분으로 분해했는데, 첫 번째 항은 χ²(μ,π)와 사전 학습된 트랜스포머의 π‑risk와 posterior regression function g_π 사이의 차이 E(e_f_T)를 곱한 형태이며, 두 번째 항은 posterior g_π와 실제 조건 기대값 E_P(Y|X) 사이의 L2 거리이다. 여기서 g_π는 사전 학습 prior에 대한 베이지안 사후 평균으로, 모든 가능한 함수 중 π‑risk를 최소화한다는 점이 강조된다.

다음 단계에서는 Theorem 2를 통해 소프트맥스 어텐션과 ReLU·GELU·SiLU와 같은 비다항 활성화 함수를 갖는 FFN 레이어를 포함한 트랜스포머가 보편적인 함수 근사 능력을 가짐을 증명한다. 특히, 모델 차원 d_model과 레이어 수를 충분히 크게 하면, 임의의 연속 함수에 대해 원하는 정밀도로 근사할 수 있음을 보이며, 이는 사전 학습된 트랜스포머가 g_π를 거의 완벽하게 학습할 수 있음을 의미한다.

섹션 4에서는 Besov 공간에 속하는 비모수 회귀 함수와 다중 인덱스 모델을 대상으로 두 가지 주요 결과를 도출한다. 첫째, posterior g_π는 사전 학습 prior의 지원 전체에 걸쳐 난이도 α에 따라 최적의 적응성을 보이며, 이는 테스트 난이도 β에 대해 자동으로 더 빠른 수렴률을 제공한다는 뜻이다. 둘째, χ² 제한 안에서 μ가 π_β와 다소 차이가 있더라도, 첫 번째 위험 항을 충분히 작은 값으로 만들 수 있기 때문에 전체 위험은 두 번째 항, 즉 posterior와 진짜 회귀 함수 사이의 차이에 의해 지배된다. 따라서 트랜스포머는 분포 이동에 강인하면서도 난이도에 따라 최적의 속도로 수렴한다.

마지막으로 Section 4.3에서는 μ에 의존하는 어떤 추정기라도 기대 위험의 수렴 속도를 g_π보다 빠르게 만들 수 없다는 하한을 제시한다. 이는 기존 연구에서 사용되는 minimax 하한보다 더 현실적인 최적성 보장을 제공한다. 실험 섹션에서는 합성 데이터와 실제 탭ular 데이터셋을 이용해 이론적 예측을 검증했으며, 난이도가 낮은 과제에서는 적은 샘플만으로도 높은 정확도를, 난이도가 높은 과제에서는 더 많은 샘플이 필요함을 확인했다. 전체적으로 논문은 사전 학습된 트랜스포머가 테스트 시점에 자동으로 난이도에 적응하고, 제한된 χ² 범위 내에서 분포 이동에 강인함을 수학적으로 증명함으로써, ICL 메커니즘에 대한 이해를 한 단계 끌어올렸다.