An Optimal Control Approach To Transformer Training

Each language version is independently generated for its own context, not a direct translation.

이 논문은 인공지능의 핵심 기술인 '트랜스포머 (Transformer)'를 훈련시키는 새로운 방법을 제안합니다. 기존 방식은 마치 미끄러운 언덕을 굴러 내려가 가장 낮은 지점을 찾는 것처럼 **경사 하강법 (Gradient Descent)**을 사용하는데, 이 방법은 종종 최적의 해가 아닌 중간에 멈추거나, 복잡한 수학적 조건이 필요하다는 한계가 있습니다.

이 논문은 이를 **최적 제어 이론 (Optimal Control Theory)**이라는 새로운 렌즈로 바라보며, "트랜스포머를 훈련한다는 것은 수많은 입자 (데이터) 를 한 번에 조종하여 목적지에 가장 잘 도달하게 하는 문제"로 정의합니다.

이 복잡한 내용을 일상적인 비유로 쉽게 설명해 드리겠습니다.

1. 핵심 아이디어: "수천 명의 군중을 한 번에 조종하기"

비유: 마술사의 군중 통제
상상해 보세요. 마술사가 수천 명의 관객 (데이터 입자) 을 한 무대에 세웠습니다. 마술사 (트랜스포머) 는 관객들이 서로 눈을 마주치며 (Self-Attention, 자기 주의 메커니즘) 서로의 위치를 파악하게 합니다.

기존 방식 (경사 하강법): 마술사가 관객 한 명 한 명을 따로따로 불러내어 "너는 조금 왼쪽으로, 너는 조금 오른쪽으로" 움직이게 합니다. 하지만 관객이 너무 많고 서로 복잡하게 얽혀 있어, 모든 사람이 완벽하게 원하는 위치로 가는지를 보장하기 어렵습니다.
이 논문의 방식 (최적 제어): 마술사는 개별 관객을 보지 않습니다. 대신 **"관객 전체의 흐름 (분포)"**을 봅니다. 마치 안개 속의 구름을 보듯, 관객들이 모여 있는 '무게의 중심'을 보고 전체를 한 번에 움직이는 명령을 내립니다. 이를 맥키 - 블라스 (McKean-Vlasov) 동역학이라고 하는데, 쉽게 말해 "개별적인 움직임이 아니라, 전체 군중의 흐름을 통제하는 것"입니다.

2. 문제 해결: "순서 잃어버린 군중을 다시 정리하기"

트랜스포머는 문장이나 시퀀스 (순서) 를 다룹니다. "사과"와 "바나나"가 순서대로 있어야 의미가 있죠. 하지만 전체 군중의 흐름을 하나로 묶어 (확률 분포로) 보면, "누가 먼저였는지"라는 순서 정보가 사라질 수 있습니다.

해결책: 이름표 (Positional Encoding) 붙이기
이 논문은 군중에게 **번호가 적힌 명찰 (위치 인코딩)**을 미리 붙여줍니다. "1 번, 2 번, 3 번..."처럼요. 그래서 전체 군중의 흐름을 하나로 묶어 분석하더라도, "아, 이 사람은 원래 1 번 자리였구나"라고 순서를 잃지 않고 유지할 수 있게 됩니다.

3. 훈련 방법: "3 단계 양자화 (Triply Quantized)"

이론적으로 완벽한 해를 찾기는 좋지만, 컴퓨터로 계산하려면 너무 복잡합니다. 무한한 가능성 중 하나를 찾아야 하니까요. 그래서 저자들은 **3 단계 양자화 (Quantization)**라는 방법을 썼습니다.

비유: 고해상도 사진을 픽셀로 나누기

1 단계 (상태 양자화): 관객들이 움직일 수 있는 공간을 아주 작은 칸 (격자) 으로 나눕니다. (예: 100x100 칸)
2 단계 (분포 양자화): 관객들이 모여 있는 '무게의 중심'도 이 작은 칸들 안에서만 계산합니다.
3 단계 (행동 양자화): 마술사가 내릴 수 있는 명령 (가중치) 도 미리 정해진 몇 가지 옵션 (예: 왼쪽, 오른쪽, 정지) 으로 제한합니다.

이렇게 하면 무한한 문제가 유한한 퍼즐이 되어, 컴퓨터가 **동적 프로그래밍 (Dynamic Programming)**이라는 방법으로 "어떤 순서로 명령을 내리면 가장 잘 될까?"를 완벽하게 계산해 낼 수 있게 됩니다.

4. 중요한 발견: "훈련은 닫힌 문, 실행은 열린 문"

이론상으로는 "현재 상태를 보고 다음 명령을 내리는 (Closed-loop)" 방식이 가장 좋습니다. 하지만 실제 트랜스포머는 훈련이 끝나면 가중치 (Weights) 를 고정하고 새로운 데이터만 입력받습니다.

논문의 통찰: 이 논문은 "훈련 과정에서 계산된 최적의 '닫힌 문' 전략은, 사실은 **훈련 데이터만 보고 미리 계산된 '열린 문' 전략 (Open-loop)**과 똑같다"고 증명했습니다.
비유: 마술사가 훈련 중에는 관객의 반응을 보고 즉흥적으로 지시를 내렸지만, 훈련이 끝난 후에는 "이 관객들이 들어오면 A, 저 관객들이 들어오면 B"라는 미리 짜인 대본만 있으면 됩니다. 이 대본이 바로 우리가 아는 '트랜스포머의 가중치'입니다.

5. 결과: "더 많은 데이터, 더 안정적인 결과"

이론적으로 증명된 바에 따르면:

최적성: 이 방법으로 찾은 해는 기존 방식보다 더 좋은 (전역 최적) 해에 가깝습니다.
강건성 (Robustness): 훈련에 사용한 데이터가 조금 변해도 (예: 다른 사람 데이터를 넣어도) 결과가 크게 흔들리지 않습니다. 마치 튼튼한 다리가 바람에 흔들리지 않는 것처럼요.
수렴: 데이터가 무한히 많아지면, 이 방법이 찾은 해는 진짜 완벽한 해에 점점 가까워집니다.

요약

이 논문은 **"트랜스포머 훈련을 개별적인 점프가 아니라, 전체 군중의 흐름을 조종하는 마술로 바꾸자"**고 제안합니다. 그리고 그 복잡한 흐름을 3 단계로 나누어 작은 퍼즐처럼 풀어서, 수학적으로 완벽에 가까운 해를 찾을 수 있음을 증명했습니다.

기존의 '경사 하강법'이 미끄러운 언덕에서 넘어지며 최적점을 놓칠 수 있다면, 이 방법은 지도와 나침반을 들고 가장 짧은 길을 계산하는 것과 같습니다. 비록 계산량이 많아 현재는 이론적 모델이나 작은 실험에 국한되지만, 트랜스포머의 작동 원리를 이해하고 더 강력한 AI 를 만드는 데 중요한 이정표가 될 것입니다.

Each language version is independently generated for its own context, not a direct translation.

1. 문제 정의 (Problem Statement)

배경: 현재 Transformer 및 대규모 언어 모델 (LLM) 은 주로 경사 하강법 (Gradient Descent) 을 기반으로 학습됩니다. 그러나 Transformer 의 손실 함수 (Loss function) 구조는 일반적으로 비볼록 (non-convex) 이고 매끄럽지 않아, 경사 하강법이 전역 최적점 (Global Optimum) 이 아닌 국소 최적점 (Local Optimum) 에 수렴할 위험이 있습니다.
핵심 과제:
1. 구조적 제약 준수: Transformer 의 실행 시 '실현된 입력 독립성 (Realized-input-independence)' (즉, 학습 후 가중치가 고정되어야 함), '앙상블 제어 (Ensemble control)' 특성 (모든 데이터 포인트에 동일한 제어 입력 적용), 그리고 '위치 의존성 (Positional dependence)'을 수학적으로 엄밀하게 반영하는 프레임워크가 필요합니다.
2. 비마르코프성 (Non-Markovian nature): Transformer 의 자기 주의 (Self-attention) 메커니즘은 각 입자 (데이터 포인트) 의 상태가 전체 앙상블의 경험적 분포 (Empirical distribution) 에 의존하므로, 입자 수준에서는 마르코프 성질이 성립하지 않습니다. 이는 동적 계획법 (Dynamic Programming) 적용을 어렵게 만듭니다.
3. 전역 최적성 보장: 볼록성이나 매끄러움에 의존하지 않고, Transformer 학습 문제에 대해 전역적으로 최적인 가중치 (제어 입력) 의 존재를 보장하고 이를 계산할 수 있는 방법이 필요합니다.

2. 방법론 (Methodology)

저자들은 Transformer 학습을 **이산 시간 제어된 입자 시스템 (Discrete-time controlled particle system)**으로 모델링하고, 이를 **확률 측도 공간 (Space of probability measures)**으로 들어올려 (Lifting) 최적 제어 문제로 재정의했습니다.

가. McKean-Vlasov 동역학 및 상태 공간 확장

입자 수준 모델링: 각 데이터 포인트를 입자로 간주하고, Transformer 블록 (피드포워드 레이어 + 단일 헤드 자기 주의) 을 상호작용하는 입자들의 동역학으로 정의합니다.
위치 인코딩 (Positional Encoding) 통합: 위치 정보를 보존하기 위해 상태 공간에 위치 인코딩을 명시적으로 포함시킵니다. 이는 측도로 들어올릴 때 순서 정보가 손실되는 것을 방지합니다.
McKean-Vlasov 동역학: 각 입자의 상태 전이가 전체 앙상블의 경험적 측도 (Empirical measure) 에 의존하는 McKean-Vlasov 유형의 동역학을 가집니다.

나. 측도 값 MDP (Measure-Valued MDP) 로의 리프팅 (Lifting)

마르코프 성질 회복: 입자 수준에서는 비마르코프적이지만, 상태 공간을 확률 측도 공간 $\mathcal{P}(\mathcal{X})$ 로 리프팅하면 시스템이 완전히 관측 가능한 마르코프 결정 과정 (MDP) 으로 변환됩니다.
전이 커널 (Transition Kernel): 측도 간의 전이는 결정론적 함수 $\Phi$ 를 통해 정의되며, 이는 약한 Feller 성질 (Weak Feller property) 을 만족합니다.
동적 계획법 (Dynamic Programming): 리프팅된 MDP 에 대해 벨만 방정식 (Bellman equations) 을 수립하여 전역 최적 정책의 존재성을 증명합니다.

다. 폐루프 (Closed-loop) 에서 개루프 (Open-loop) 로의 변환

실제 구현과의 호환성: 최적 제어 이론에서 얻은 폐루프 정책 (현재 상태에 기반한 피드백) 은 Transformer 의 실제 작동 방식 (학습 후 가중치 고정) 과 다릅니다.
동치성 증명: 리프팅된 문제의 결정론적 폐루프 정책은 **초기 분포에 의존하는 개루프 정책 (Initial-distribution dependent open-loop policy)**과 동치임을 증명합니다. 이는 학습이 완료된 후 Transformer 의 가중치를 고정하는 표준적인 방식과 정확히 일치합니다.

라. 삼중 양자화 학습 알고리즘 (Triply Quantized Training Scheme)

계산 가능성 확보: 연속적인 상태, 측도, 제어 공간을 직접 다루는 것은 계산적으로 불가능하므로, 세 가지 차원을 양자화 (Quantization) 합니다.
1. 상태 공간 양자화: 입자 상태를 유한 그리드로 근사.
2. 측도 공간 양자화: 확률 측도를 유한한 reconstruction points 집합으로 근사 (Reznik 의 방법 활용).
3. 제어 공간 (행렬) 양자화: 가중치 행렬을 유한 집합으로 근사.
근사 최적성: 양자화된 모델에서 얻은 최적 정책이 원래 문제의 $\epsilon$ -근사 최적 (Near-optimal) 해임을 증명합니다.

3. 주요 기여 (Key Contributions)

Transformer 의 엄밀한 최적 제어 공식화: Transformer 아키텍처를 공유 제어 하의 McKean-Vlasov 동역학 시스템으로 모델링하고, 위치 인코딩을 포함한 측도 값 MDP 로 확장하여 학습 문제를 재정의했습니다.
전역 최적 정책의 존재성 증명: 컴팩트성 가정 하에서 동적 계획법을 통해 전역 최적 제어 (가중치) 의 존재를 수학적으로 증명했습니다. 이는 경사 하강법의 국소 최적점 한계를 극복합니다.
이론적 최적 정책과 실제 학습의 연결: 리프팅된 MDP 의 폐루프 정책이 Transformer 의 실행 시나리오에 부합하는 '초기 데이터 의존적 개루프 정책'으로 변환될 수 있음을 보였습니다. 이는 학습된 가중치가 고정된다는 표준 관행과 이론적으로 일치함을 의미합니다.
삼중 양자화 기반 학습 알고리즘: 상태, 측도, 행동을 모두 양자화하여 계산 가능한 유한 상태 MDP 를 구성하고, 이 알고리즘이 원래 문제에 대해 점근적으로 최적에 수렴함을 보였습니다.
강건성 및 일반화 이론: 학습 데이터의 분포가 실제 분포에 수렴할 때, 최적 제어 행동도 실제 분포에 대해 최적에 수렴함을 증명했습니다 (약한 수렴 및 $\Gamma$ -수렴). 이는 일반화 문제 (Generalization problem) 에 대한 최적 제어적 해석을 제공합니다.

4. 결과 (Results)

이론적 결과:
- 리프팅된 MDP 에 대해 동적 계획법 방정식이 잘 정의되며, 전역 최적 정책이 존재함을 보였습니다.
- 삼중 양자화 모델의 최적 정책이 원래 문제의 $\epsilon$ -근사 최적임을 증명했습니다.
- 가치 함수 (Value function) 가 초기 경험적 측도의 섭동에 대해 연속적임을 보였으며, 이는 데이터가 증가함에 따라 학습된 정책이 안정적이고 일관됨을 의미합니다.
실험적 결과 (Toy Problem):
- 단순한 자기 주의 (Self-attention) 계층을 근사하는 toy problem 에서 삼중 양자화 알고리즘을 적용했습니다.
- 양자화 수준 증가에 따른 성능 향상: 행동 (Action) 양자화 수준을 높임에 따라 훈련 오차와 테스트 오차가 감소하는 것을 확인했습니다.
- 수렴성: 양자화 수준이 세밀해질수록 오차가 감소하며 최적 해에 근접함을 보여주었습니다.
- 계산 비용: 실행 시간은 행동 수의 제곱 ( $O(M^2)$ ) 에 비례하여 증가하는 것으로 관찰되었습니다.

5. 의의 및 의의 (Significance)

새로운 학습 패러다임 제시: Transformer 학습을 경사 하강법 (Gradient Descent) 에만 의존하지 않고, 최적 제어 이론과 동적 계획법을 기반으로 한 대안적 접근법을 제시했습니다.
구조적 이해 증진: Transformer 의 비선형적이고 비볼록한 구조를 McKean-Vlasov 동역학 및 앙상블 제어 관점에서 해석함으로써, 모델의 내부 작동 원리에 대한 구조적 이해를 심화시켰습니다.
이론적 보장: 볼록성이나 미분 가능성과 같은 강한 가정이 없어도 전역 최적해의 존재를 보장할 수 있음을 보여주어, Transformer 와 같은 복잡한 신경망의 이론적 기반을 강화했습니다.
일반화 문제 해결: 학습 데이터의 분포 변화에 대한 가치 함수의 연속성을 증명함으로써, 학습된 모델의 일반화 성능에 대한 이론적 근거를 제공했습니다.

결론적으로, 이 논문은 Transformer 학습을 최적 제어 문제로 재정의하고, 이를 통해 전역 최적성, 강건성, 그리고 실제 구현 (가중치 고정) 과의 호환성을 모두 만족하는 이론적 프레임워크를 제시한 획기적인 연구입니다. 비록 현재 제안된 양자화 알고리즘이 대규모 데이터에 대한 확장성 (Scalability) 면에서는 제한적이지만, Transformer 의 구조와 최적성 원리에 대한 깊은 통찰을 제공한다는 점에서 중요한 의의를 가집니다.