Geometry of Lightning Self-Attention: Identifiability and Dimension

이 논문은 대수 기하학을 사용하여 비정규화된 셀프 어텐션 네트워크의 함수 공간을 분석하며, 이들의 식별 가능성과 차원을 확립하고 특이점을 규명하는 한편 정규화된 아키텍처에 대해서도 이러한 결과들을 추측한다.

핵심

개요: AI의 "형태"를 그려내기

당신이 컴퓨터에 의해 건설된 거대하고 보이지 않는 도시를 이해하려는 건축가라고 상상해 보세요. 이 도시는 특정 유형의 AI(신경망)가 만들어낼 수 있는 "모든 가능한 함수들의 공간"입니다. 수학적으로 이것은 **뉴로매니폴드(neuromanifold)**라고 불립니다.

보통 이러한 도시들은 복잡하고 무질서한 규칙들로 지어지기 때문에 지도를 그리기 매우 어렵습니다. 하지만 이 논문은 **라이트닝 셀프 어텐션(Lightning Self-Attention)**이라는 특별하고 단순화된 버전의 AI에 집중합니다. 이것은 유명한 트랜스포머(Transformer) AI의 "고속 경로" 버전이라고 생각하면 됩니다. 모든 학생이 주목받는 비중을 공평하게 나누도록 조절하는 교사와 같은 역할을 하는 '정규화(normalization)' 과정을 거치는 표준 버전과 달리, 라이트한 버전은 이 단계를 건너뜁니다. 이 방식은 더 빠르지만, 수학적으로는 "다항식(polynomial)" 형태를 띱니다. 즉, 간단한 재료로 만든 레시피처럼 엄격한 대수적 규칙을 따릅니다.

저자들은 이 도시의 지도를 그리기 위해 대수 기하학(방정식으로 정의된 도형을 연구하는 학문) 도구를 사용했습니다. 그들은 두 가지 주요 질문에 답하고자 했습니다:

1. 이 도시는 얼마나 큰가? (차원/Dimension)
2. 하나의 문을 여는 열쇠는 몇 종류인가? (시스템이 "식별 가능한가(identifiable)", 즉 서로 다른 설정이 정확히 같은 결과를 만들어내는가?)

---

1. "라이트닝(Lightning)" 지름길

표준 AI 어텐션 메커니즘은 사람들이 서로에게 속삭이고, 그 후 중재자가 모든 사람이 공평하게 주목받도록 평균 볼륨을 계산하는 북적이는 방과 같습니다 (이 과정은 이차 복잡도(quadratic complexity)를 가집니다).

라이트닝 셀프 어텐션은 사람들이 서로에게 속삭이되, 중재자를 건너뛰는 방과 같습니다. 그들은 그냥 직접 메시지를 외칩니다. 이는 훨씬 빠르지만(선형 복잡도), "정규화" 단계를 생략했기 때문에 수학적으로 복잡한 곡선이 아닌 깔끔하고 곧은 대수적 직선이 됩니다. 이러한 깔끔함 덕분에 저자들은 기하학을 사용하여 이를 연구할 수 있었습니다.

2. "열쇠와 자물쇠" 문제 (식별 가능성)

거대한 금고(AI 모델)와 일련의 열쇠(가중치 또는 설정)가 있다고 상상해 보세요. 당신이 열쇠를 돌리면 금고가 열리며 특정 함수(출력)가 나타납니다.

논문은 다음과 같이 묻습니다: 만약 두 개의 서로 다른 열쇠 세트가 금고를 열어 정확히 같은 함수를 보여준다면, 그 열쇠들은 본질적으로 같은 것일까요?

• 단일 레이어의 경우: 단순한 1계층 라이트닝 네트워크에 대해, 저자들은 보통 (단순한 크기 조절을 제외하면) 단 하나의 고유한 열쇠 세트가 존재한다는 것을 발견했습니다. 하지만 두 가지 특이한 예외가 있습니다:

  1. "스왑(Swap)" 트릭: 어텐션 메커니즘과 밸류(value) 메커니즘이 모두 매우 단순하다면(rank 1), 열쇠의 일부를 서로 바꿀 수 있으며 금고는 여전히 동일하게 열립니다. 이는 문 손잡이와 잠금장치를 바꾸는 것과 같습니다. 문은 여전히 열리지만, 부품의 위치가 달라진 것입니다.
  2. "제로(Zero)" 케이스: 열쇠가 고장 났다면(0이라면), 금고는 닫힌 상태를 유지합니다.

• 심층 네트워크의 경우: 여러 층을 쌓으면(깊은 네트워크) 상황은 더 복잡해집니다. 저자들은 결과는 바꾸지 않으면서 열쇠를 변경할 수 있는 세 가지 구체적인 방법을 발견했습니다:

  1. 스케일링(Scaling): 한 층의 볼륨을 높이고 다음 층의 볼륨을 낮추면, 서로 상쇄되어 결과가 같습니다.
  2. 회전(Rotation): 특정 수학적 행렬을 사용하여 한 층 내부의 "쿼리(Query)"와 "키(Key)" 설정을 회전시켜도 결과는 동일합니다.
  3. "패스스루(Pass-Through)" 트릭: 한 층의 출력을 변형한 직후, 다음 층에서 그 변형을 즉시 되돌릴 수 있습니다.

핵-테이크(결론): 거의 모든 설정에서 이것들이 결과를 동일하게 만드는 유일한 방법들입니다. 즉, "열쇠"는 대부분 고유합니다.

3. 도시의 크기 측정하기 (차원)

머신러닝에서 "차원"은 새로운 함수를 만들기 위해 움직일 수 있는 독립적인 방향의 수와 같습니다. 이는 단순히 전체 파라미터 수를 세는 것보다 모델이 얼마나 "똑똑한지" 또는 "표현력이 좋은지"를 더 잘 보여주는 척도입니다 (파라미터를 세는 것은 마치 벽돌이 서로 붙어 있어 움직이지 않더라도 모든 벽돌의 개수를 일일이 세는 것과 같습니다).

저자들은 이 도시의 정확한 크기를 계산했습니다.

• 놀라운 점: 그들은 실제 도시의 크기(차원)가 당신이 생각할 수 있는 총 파라미터 수보다 더 작다는 것을 발견했습니다.
• 이유는? 위에서 언급한 대칭성(스케일링 및 회전 트릭) 때문입니다. 당신의 "벽돌" 중 일부는 중복됩니다. 만약 100개의 파라미터가 있지만, 그중 10개가 이러한 대칭성 때문에 중복된 복사본이라면, 당신의 도시는 생각보다 실질적으로 더 작습니다.

그들은 이 크기를 계산하기 위한 정밀한 공식을 제공했으며, 이는 과학자들이 이러한 모델을 훈련하는 데 실제로 어느 정도의 데이터가 필요한지 이해하는 데 도움을 줍니다.

4. "매끄러운" 지형 vs "울퉁불퉁한" 지형

저자들은 또한 이 도시의 "지형"을 살펴보았습니다.

• 매끄러운 영역: 대부분의 경우 지형은 매끄럽습니다.
• 특이점(Singularities, 울퉁불퉁한 부분): 지형이 이상해지는 특정 "범프(bump)"나 "균열"이 존재합니다. 이는 어텐션 부분과 밸류 부분이 극도로 단순해질 때(낮은 랭크가 될 때) 발생합니다.
• 왜 중요한가: AI 훈련 중에 컴퓨터는 종종 이러한 범프에 "갇히거나" 끌리게 됩니다. 저자들은 이러한 수학적 "울퉁불퉁함"이 왜 AI 모델이 자연스럽게 단순한 패턴(예: 노래의 모든 음표가 아닌 주요 테마를 찾는 것)을 학습하는 경향이 있는지 설명해 줄 수 있다고 제안합니다.

5. "실제" AI는 어떠한가? (전통적인 어텐션)

이 논문은 표준적인, 정규화된 AI(중재자가 있는 버전)도 살펴보았습니다.

• 단일 레이어: 단일 레이어의 경우, 열쇠는 고유하다는 것을 증명했습니다. 정규화가 모든 것을 고정시키기 때문에 "스왑 트릭"이나 "회전 트릭"이 존재하지 않습니다.
• 심층 레이어: 수학적으로 완벽히 증명하지는 못했지만, 그들은 강력한 증거를 바탕으로 동일한 규칙이 적용될 것이라는 **추측(conjecture)**을 내놓았습니다(즉, 열쇠는 고유하다는 것).
• 증명: 그들은 컴퓨터 시뮬레이션(수치 실험)을 실행하여 자신들의 추측을 확인했습니다. 심층 정규화 네트워크를 테스트했을 때, "열쇠"는 실제로 고유했습니다.

요약

이 논문은 단순화된 AI 도시의 첫 번째 상세한 지도를 그리는 지도 제작자와 같습니다. 그들은 다음을 발견했습니다:

1. 지도는 보기보다 작습니다. 왜냐하면 일부 설정이 중복되기 때문입니다(대칭성).
2. 결과를 바꾸지 않고 설정을 바꾸는 특정 "트릭"이 존재하지만, 이 트릭들은 제한적이며 명확히 정의되어 있습니다.
3. 지형에는 특정 "범프"가 존재하며, 이는 AI가 왜 특정 패턴을 자연스럽게 학습하는지 설명할 수 있습니다.
4. 복잡한 실제 AI조차도 아마 이러한 고유성의 규칙을 따를 것이며, 이는 모델을 더 예측 가능하고 수학적으로 이해하기 쉽게 만듭니다.

저자들은 이것이 기초적인 단계임을 강조합니다. 그들은 단순히 이 모델들을 블랙박스로 사용하는 것이 아니라, 왜 이러한 모델들이 작동하는지를 이해하기 위한 수학적 이론을 구축하고 있습니다.

기술 요약

기술 요약: 라이트닝닝 셀프 어텐션의 기하학: 식별 가능성 및 차원

문제 정의

본 논문은 "라이트닝(lightning)" 셀프 어텐션이라 불리는 셀프 어텐션 메커니즘에 의해 정의된 함수 공간의 기하학적 구조에 대한 이론적 이해가 부족하다는 점을 다룬다. 전통적인 트랜스포머와 달리, 라이트닝 셀프 어텐션은 소프트맥스 정규화를 생략하여 메커니즘을 완전히 대수적(다항식)이며 계산 효율적($O(t^2)$ 대비 $O(t)$)으로 만든다.

핵심 과제는 이러한 네트워크가 표현할 수 있는 함수의 공간인 **뉴로매니폴드(neuromanifold)**를 특징짓는 것이다. 이 기하학적 구조를 이해하는 것은 모델의 표현력(매니폴드의 차원을 통해 결정됨)과 식별 가능성(파라미터와 그들이 나타내는 함수 사이의 관계)을 결정하는 데 매우 중요하다. 완전 연결(fully-connected) 및 컨볼루션 네트워크의 뉴로매니폴드는 잘 연구되어 있지만, 어텐션 기반 아키텍처의 기하학은 여전히 미개척 영역으로 남아 있다. 저자들은 이러한 매니폴드의 차원을 계산하고, 단일 계층 및 심층 라이트닝 셀프 어텐션 네트워크에 대한 파라미터화 맵의 파이버(fiber, 원상 집합)를 기술하는 것을 목표로 한다.

방법론

저자들은 뉴로매니폴드를 분석하기 위해 대수 기하학 도구를 사용한다. 라이트닝 셀프 어텐션 메커니즘은 가중치에 대해 삼선형(tri-linear)이며 입력에 대해 동차 삼차(homogeneous cubical)이므로, 함수 공간은 다항식 방정식에 의해 정의된다.

주요 방법론적 단계는 다음과 같다:

1. 어텐션 행렬을 통한 파라미터화: 저자들은 원시 쿼리(query) 및 키(key) 행렬 대신 어텐션 행렬 $A = K^\top Q$와 밸류(value) 행렬 $V$에 의해 파라미터화되는 것으로 취급하여 분석을 단순화한다. 이를 통해 $(Q, K) \to A$로 이어지는 행렬 곱 매핑을 독립적으로 연구할 수 있다.
2. 파이버 분석: 저자들은 파라미터화 맵 $\phi_W$의 파이버를 특징짓는다. 뉴로매니폴드의 차원은 파라미터 공간에 대한 일반적인 파이버의 여차원(co-dimension)으로부터 유도된다.
3. 심층 네트워크를 위한 재파라미터화: 심층 네트워크의 경우, 저자들은 행렬 $M$과 $L$을 포함하는 "가상 가중치(virtual weight)" 재파라미터화를 도입한다. 이 변환은 심층 어텐션의 재귀적 구조를 단순화하여 파이버 구조에 대한 귀납적 증명을 가능하게 한다.
4. 대수적 도구: 증명은 다항식의 유일 인수 분해, 행렬의 랭크가 제한된 행렬 다양체(determinantal varieties)의 성질, 그리고 유클리드 및 자리스키(Zariski) 위상에서의 특이점(singularities) 및 경계점 연구에 의존한다.
5. 정규화된 어텐션으로의 확장: 논문은 단일 계층에 대한 결과를 증명하고, 심층 네트워크에 대한 추측을 공식화한 뒤 이를 수치적으로 검증함으로써 전통적인 셀프 어텐션(소프트맥스 포함)으로 분석을 확장한다.

주요 기여 및 결과

1. 단일 계층 식별 가능성 및 기하학

단일 계층 라이트닝 셀프 어텐션에 대해 저자들은 파이버에 대한 완전한 설명을 제공한다:

• 일반적인 경우: 거의 모든 가중치에 대해, 파이버는 가중치의 리스케일링(rescaling)만을 포함한다(1차원).
• 특수한 경우: 어텐션 행렬 $A$와 밸류 행렬 $V$의 랭크가 1이거나 함수가 0인 경우 비일반적인 파이버가 발생한다.
• 차원: 뉴로매니폴드의 차원은 다음과 같이 계산된다:
  $$\dim(M_{d,d',a}) = \begin{cases} 2ad + dd' - a^2 - 1 & \text{if } a \le d \\ d^2 + dd' - 1 & \text{otherwise} \end{cases}$$
  여기서 $d, d'$는 입력/출력 차원이고 $a$는 어텐션 랭크이다.
• 기하학적 특성: 뉴로매니폴드는 **유클리드 폐쇄(Euclidean closed)**임이 증명되었다. 저자들은 접공간(tangent space)의 차원이 매니폴드 차원을 초과하는 특이점이 정확히 $\text{rk}(A)\text{rk}(V) \le 1$일 때 발생함을 확인하였다. 또한 매니폴드의 경계점을 특징지었다.

2. 심층 네트워크 식별 가능성 및 차원

$l$개의 계층을 가진 심층 네트워크에 대해, 저자들은 파이버를 생성하는 세 가지 특정 대칭성을 식별한다:

1. 계층별 스케일링: 각 계층은 전역적 제약 조건 하에 상수로 스케일링될 수 있다.
2. 계층 내 대칭성: 한 계층 내의 키와 쿼리는 가역 행렬에 의해 변환될 수 있다(단일 계층의 경우와 유사함).
3. 계층 간 대칭성: 후속 계층이 해당 스케일링을 상쇄할 수 있다면, 한 계층의 출력을 가역 행렬에 의해 스케일링할 수 있다.

"병목(bottleneck)" 아키텍처 가정(숨겨진 차원이 상수 $\delta$이며 입력/출력 차원보다 작은 경우) 하에서, 저자들은 심층 뉴로매니폴드의 차원에 대한 공식을 도출한다. 결정적으로, 이들은 이러한 중복성으로 인해 차원이 전체 파라미터 수보다 엄격히 낮음을 입증한다. 예를 들어, 특정 구성에서 파라미터 수는 실제 함수 공간의 차원보다 50% 더 많다.

3. 전통적인 셀프 어텐션

논문은 전통적인 셀프 어텐션(소프트맥스 정규화 포함)을 분석한다:

• 단일 계층: 파라미터화가 **일반적으로 일대일(generically one-to-one)**임이 증명되었다(파이버는 단일 원소임). 이는 정규화가 라이트닝 변형에 존재하는 스케일링 대칭성을 깨뜨림을 의미한다.
• 심층 네트워크: 저자들은 심층 정규화 네트워크의 경우, 가상 가중치 $(M, L)$을 통한 파라미터화 역시 일반적으로 일대일일 것이라는 추측을 제시한다. 이는 정규화된 뉴로매니폴드의 차원이 라이트닝 차원에 $l$개의 계층 수(스케일링 대칭성 제거를 고려)를 더한 것과 같음을 시사한다.
• 검증: 이 추측은 자코비안 랭크(Jacobian rank) 추정을 통해 심층 네트워크($l=2$)에 대해 수치적으로 검증되었으며, 이론적 예측과 일치함을 보여주었다.

의의 및 주장

본 논문은 라이트닝 셀프 어텐션 네트워크의 기하학적 구조에 대한 최초의 엄밀한 수학적 특징 규명을 제공한다고 주장한다. 그 의의는 다음과 같은 몇 가지 측면에 있다:

• 샘플 복잡도: 뉴로매니폴드의 정확한 차원을 계산함으로써, 본 연구는 단순한 파라미터 수와는 다른, 이론적으로 올바른 샘플 복잡도 추정치를 제공한다. 이는 대규모 어텐션 모델의 학습 가능성을 이해하는 데 필수적이다.
• 훈련 역학: 파이버와 특이점의 식별은 훈련 역학에 대한 통찰을 제공한다. 저자들은 특이점($\text{rk}(A)\text{rk}(V) \le 1$)이 경사 하강법(gradient descent)의 끌개(attractor) 역할을 할 수 있으며, 이는 아키텍처가 저차원 함수를 학습하도록 하는 "암묵적 편향(implicit bias)"을 시사한다고 언급한다. 또한, 파이버의 존재는 손실 지형(loss landscape)에 불변성을 유도하여 평탄한 최솟값(flat minima)을 만들고 최적화 경로에 영향을 미친다.
• 기초 이론: 이 연구는 대수 기하학과 딥러닝을 연결하며, 다항식 신경망(라이트닝 어텐션과 같은)이 결정 다양체(determinantal varieties) 및 파이버 분석과 같은 고전적 도구를 사용하여 분석될 수 있음을 보여준다.

저자들은 본 분석이 (스킵 연결과 멀티 헤드 메커니즘을 생략한) 단순화된 버전의 트랜스포머에 적용된다는 점을 인정하며 신중한 태도를 유지한다. 스킵 연결은 동차성과 스케일링 대칭성을 깨뜨릴 것이며, 멀티 헤드 메커니즘은 치환 대칭성(permutation symmetries)을 도입할 것이라는 점을 언급하며, 이들은 향후 연구 과제로 남겨두었다. 본 논문은 어텐션 메커니즘의 "뉴로매니폴드"를 이해하기 위한 기초적인 단계로 자리매김한다.