Toward a First-Principles Update Geometry for the Language-Model Head

현대 인공지능의 광활한 풍경 속에서, 가장 눈에 띄는 성취는 종종 인간과 유사한 텍스트를 생성하는 시스템으로부터 옵니다. 이러한 대화의 이면에는 방대한 양의 데이터로 학습된 디지털 뇌인 복잡한 머신러닝 모델이 자리 잡고 있습니다. 이 기계의 처리 체계 맨 끝에는 언어 모델 헤드(language-model head)라고 알려진 특정 구성 요소가 위치합니다. 이 구성 요소의 역할은 매우 단순해 보입니다. 컴퓨터가 생각하고 있는 바에 대한 수학적 표현인 은닉 내부 상태(hidden internal state)를 받아들여, 이를 전체 어휘집에 있는 모든 단어에 대한 확률 목록으로 변환하는 것입니다. 만약 모델이 "The sky is"라는 문장을 완성하려 한다면, 이 구성 요소는 다음 단어가 "blue", "cloudy", 또는 "falling"일 확률이 각각 얼마인지 계산하여 각 가능성에 점수를 부여합니다. 수년 동안 엔지니어들은 이 최종 변환을 구동하는 수학을 표준적인 문제로 취급해 왔으며, 네트워크의 초기 레이어에 적용하는 것과 동일한 일반적인 규칙을 적용하여 모델의 가중치를 조정해 왔습니다.

하지만 새로운 관점은 이 마지막 단계가 기계의 나머지 부분과는 근본적으로 다르다고 제안합니다. 출력값이 확률 분포, 즉 반드시 100%가 되어야 하는 백분율의 집합이기 때문에, 이를 변화시키는 규칙은 단순히 숫자를 바꾸는 규칙이 아니라 확률의 고유한 기하학적 구조를 존중해야 한다는 것입니다. 컬럼비아 대학교의 아디티아 소마순다람(Aditya Somasundaram)의 최근 논문은 이 마지막 레이어와 확률 변환을 하나의 통합된 모듈로 취급함으로써 이 아이디어를 탐구합니다. 저자는 이 부분의 개선 방법을 이해하기 위해서는 단순히 숫자가 얼마나 변하는지가 아니라, 업데이트가 단어들 사이의 상대적 오즈(relative odds)를 어떻게 변화시키는지 살펴보아야 한다고 주장합니다. 이 연구는 두 상대적 오즈 사이의 변화를 측정하는 특정 수학적 거리를 사용하여 업데이트의 크기를 측정하는 새로운 방법을 도출합니다. 그 결과, 변화의 "크기"가 서로 다른 단어들에 대한 모델의 내부 표현의 확산(spread)에 의해 결정되는 새로운 기하학적 그림이 그려지며, 이는 모델이 더 효율적으로 학습할 수 있게 하는 새로운 종류의 옵티마이저(optimizer)에 대한 제안으로 이어집니다.

이 조사의 핵심은 변화를 어떻게 측정할 것인가라는 질문에서 시작됩니다. 컴퓨터가 실수를 통해 학습하기 위해 내부 설정을 업데이트할 때, 이는 거대한 숫자 테이블에 아주 미세한 조정을 가하는 과정입니다. 표준적인 훈련 방법에서 엔지니어들은 흔-히 이 조정의 크기를 확률로 변환되기 전의 원시 숫자(raw numbers)에 가해질 수 있는 최대 변화량으로 측정합니다. 그러나 저자는 마지막 레이어의 경우, 원시 숫자가 최종 산물이 아니라 확률이 최종 산물이라는 점을 지적합니다. 소프트맥스(softmax)라고 알려진 변환 과정은 특별한 성질을 가지고 있습니다. 만약 리스트의 모든 숫자에 동일한 양을 더한다면, 결과적인 확률은 전혀 변하지 않는다는 것입니다. 이는 원시 숫자의 크기를 측정하는 것이 오도될 수 있음을 의미하는데, 왜냐하면 그것은 최종 출력에 아무런 영향을 미치지 않는 변화까지 계산하기 때문입니다. 이를 해결하기 위해, 이 논문은 힐베르트 사영 거리(Hilbert's projective distance)라는 개념을 활용합니다. 이것은 두 확률 집합 사이의 거리를 측정하는 방식으로, 어떤 두 단어 사이의 비율이 어떻게 변하는지에 전적으로 집중합니다. 이는 숫자의 절대적인 크기는 무시하고, 한 단어가 다른 단어에 대해 갖는 상대적인 입지만을 살펴봅니다.

이 특정한 거리 측정법을 적용함으로써, 연구자는 업데이트의 기하학적 구조와 모델의 동작 사이의 놀랍고도 명쾌한 관계를 발견했습니다. 연구는 업데이트가 확률 분포에 일으킬 수 있는 최대 변화가 업데이트 행렬 내 행(row)들의 물리적 확산과 직접적으로 연결되어 있음을 보여줍니다. 업데이트를 각 단어에 대응하는 벡터들의 집합이라고 상상해 보십시오. 확률을 흔들어 놓을 수 있는 업데이트의 "크기"는 가장 멀리 떨어져 있는 두 벡터 사이의 거리에 의해 결정됩니다. 만약 서로 다른 단어들을 위한 벡터들이 서로 밀집되어 있다면, 업데이트는 작고 안전합니다. 만약 벡터들이 멀리 퍼져 있다면, 업데이트는 크며 모델의 예측에 격렬한 변동을 일으킬 수 있습니다. 이 발견은 언어 모델 헤드를 업데이트하는 문제를 재정립합니다. 즉, 숫자의 전체적인 크기에 걱정하는 대신, 단어들을 나타내는 점들의 구름(cloud of points)의 직경을 관리하는 것이 목표가 됩니다.

이러한 기하학적 통찰은 최근 신경망의 다른 부분에서 성공을 거둔 뮤온(Muon)이라는 방법에서 영감을 얻어, 이러한 업데이트를 구성하는 방법에 대한 새로운 제안으로 이어집니다. 뮤온은 업데이트의 서로 다른 방향들이 모두 동일하게 강력하도록 균형을 맞춤으로써, 모델이 학습 지형의 좁은 골짜기에 갇히는 것을 방지합니다. 저자는 이와 유사한 원리가 언어 모델 헤드에도 적용되어야 하지만, 약간의 차이가 있어야 한다고 제안합니다. 단일 방향의 강도를 조절하는 대신, 목표는 모든 단어 쌍 사이의 거리를 가능한 한 동일하게 만드는 것이어야 합니다. 이상적인 업데이트는 모든 단어 벡터를 최대한 균등하게 퍼뜨려, 모든 단어가 서로 거의 같은 거리에 있게 함으로써 완벽하게 균형 잡힌 구름을 만드는 것입니다. 이는 모델이 어떤 단어 쌍의 구별도 다른 쌍만큼 민감하게 다루도록 보장할 것입니다.

그러나 이 논문은 또한 이러한 이상적인 상태에 존재하는 물리적인 한계를 식별합니다. 언어 모델 헤드에서는 어휘집의 단어 수가 이를 표현하기 위해 사용 가능한 차원보다 훨씬 더 많습니다. 수많은 점을 좁은 공간 안에 배치하여 모두가 정확히 같은 거리에 있게 만드는 것은 수학적으로 불가능합니다. 평평한 종이 위에 백 개의 점을 모두 등거리가 되도록 배치할 수 없는 것처럼, 어휘가 방대하고 은닉 공간이 작은 경우 단어 벡터들을 완벽하게 등거리로 만들 수는 없습니다. 연구는 이러한 제약을 인정하며, 목표는 최선의 근사치를 찾는 것이 되어야 한다고 제안합니다. 제안된 옵티마이저는 어떤 두 단어 사이의 최소 거리를 최대화하는 동시에, 최대 거리는 안전한 범위 내로 유지하려고 노력할 것입니다. 이 접근 방식은 업데이트가 허용하는 기하학적 구조 내에서 최대한 고르게 분포되도록 하여, 특정 단어 쌍이 무시되거나 구별 불가능해지는 동시에 다른 단어들이 너무 멀리 밀려나는 일이 없도록 하는 것을 목표로 합니다.

이 연구의 함의는 주로 이론적이고 구조적이며, 언어 생성의 마지막 단계를 바라보는 새로운 렌즈를 제공합니다. 저자는 이 방법이 실제로 더 효과적임을 증명하는 완전히 훈련된 모델을 구축했다고 주장하는 것이 아니라, 확률의 제1 원칙을 따른다면 업데이트의 기하학적 구조가 어떠해야 하는지에 대한 엄밀한 유도를 제공하는 것입니다. 논문은 이 마지막 레이어를 표준적인 최적화 기법에 그대로 맡려두는 현재의 방식들이, 해당 모듈의 특수한 기능을 존중할 기회를 놓치고 있다고 주장합니다. 확률로의 변환을 업데이트 과정의 필수적인 부분으로 취급하고, 소프트맥스 함수의 불변성을 존중하는 거리 측정법을 사용함으로써, 제안된 기하학적 구조는 학습 지형을 항해하는 더 자연스러운 방법을 제시합니다. 연구는 완벽한 등거리 배치가 거대한 어휘집에 대해서는 불가능할지라도, 근사적인 등거리 구성을 추구하는 것이 언어 모델 헤드의 독특한 요구 사항에 특화된 미래의 옵티마이저를 설계하기 위한 명확하고 원칙적인 방향을 제공한다고 결론짓습니다.