28.25 아인슈타인 합산을 활용한 텐서 연산의 통합적 표현
아인슈타인 합산(Einstein summation) 규약은 텐서 지표의 반복이 나타나면 해당 지표에 대해 암묵적으로 합산이 수행된다는 표기법이며, 원래 일반 상대성 이론의 텐서 계산을 간결하게 기술하기 위해 아인슈타인이 도입하였다. 이 규약은 현대 수치 계산 환경에서 einsum이라는 이름의 함수로 재탄생하여, 원소별 곱셈, 축약, 축소, 전치, 내적, 외적 등 다수의 텐서 연산을 단일 인터페이스로 표현할 수 있게 한다. 본 절에서는 아인슈타인 규약이 어떻게 텐서 연산을 통합적으로 표현할 수 있는지 다룬다.
1. 아인슈타인 합산 규약의 기본 형식
텐서의 성분 표현에서 동일한 지표가 두 번 나타나면 그 지표에 대한 합산이 암묵적으로 수행된다. 예를 들어 다음 식에서 지표 k는 좌변에 나타나지 않고 우변의 두 인자에 반복되므로 합산 지표(dummy index)로 해석된다.
C_{ij} = A_{ik} B_{kj} \equiv \sum_{k} A_{ik} B_{kj}
좌변에 나타나는 지표는 자유 지표(free index)로, 출력 텐서의 차원을 구성한다. einsum 함수는 이 규약을 문자열 형태의 방정식으로 표현하여 입력 지표와 출력 지표를 명시한다. 예를 들어 행렬-행렬 곱은 ik,kj->ij로 표현된다. 여기서 -> 좌측은 입력 텐서들의 지표이고, 우측은 출력 텐서의 지표이다.
2. 통합적 표현의 체계
einsum은 다음과 같은 일반적 계산 절차로 정의될 수 있다. 입력 텐서들의 지표 집합의 합집합을 \mathcal{I}라 하고, 출력 지표 집합을 \mathcal{O} \subseteq \mathcal{I}라 하자. 축소될 지표 집합은 \mathcal{R} = \mathcal{I} \setminus \mathcal{O}이다. 그러면 출력 텐서의 원소는 다음과 같이 계산된다.
C_{\mathcal{O}} = \sum_{\mathcal{R}}\ \prod_{t=1}^{T} A^{(t)}_{\mathcal{I}^{(t)}}
여기서 T는 입력 텐서의 개수이고, A^{(t)}는 t번째 입력 텐서이다. 이 일반 공식은 모든 다중선형 연산을 단일 수식으로 포괄한다.
3. 대표적 사례
다음 표는 einsum 표기로 표현된 주요 텐서 연산을 나열한다.
| 연산 | einsum 표기 | 수식 |
|---|---|---|
| 항등 | ij->ij | C_{ij} = A_{ij} |
| 전치 | ij->ji | C_{ji} = A_{ij} |
| 트레이스 | ii-> | c = \sum_i A_{ii} |
| 대각선 추출 | ii->i | c_i = A_{ii} |
| 열 합 | ij->j | c_j = \sum_i A_{ij} |
| 행 합 | ij->i | c_i = \sum_j A_{ij} |
| 내적 | i,i-> | c = \sum_i a_i b_i |
| 외적 | i,j->ij | C_{ij} = a_i b_j |
| 행렬-벡터 곱 | ij,j->i | c_i = \sum_j A_{ij} x_j |
| 행렬-행렬 곱 | ij,jk->ik | C_{ik} = \sum_j A_{ij} B_{jk} |
| 배치 행렬 곱 | bij,bjk->bik | C_{bik} = \sum_j A_{bij} B_{bjk} |
| 아다마르 곱 | ij,ij->ij | C_{ij} = A_{ij} B_{ij} |
| 이중 축약 | ijkl,kl->ij | C_{ij} = \sum_{k,l} A_{ijkl} B_{kl} |
| 어텐션 점수 | blhd,bmhd->blhm | C_{b\ell h m} = \sum_d Q_{b\ell hd} K_{bmhd} |
위 표가 보여주듯, einsum은 단일 표기법으로 매우 광범위한 텐서 연산을 기술할 수 있다.
4. 딥러닝에서의 활용 예
트랜스포머의 다중 헤드 어텐션에서 질의 Q, 키 K, 값 V 텐서가 각각 형상 (B, L, H, d_k)를 가질 때, 어텐션 점수와 가중합은 다음과 같이 표현된다.
S_{b\ell h m} = \sum_{d=0}^{d_k - 1} Q_{b\ell h d}\, K_{b m h d}
O_{b\ell h d} = \sum_{m=0}^{L - 1} A_{b\ell h m}\, V_{b m h d}
이러한 연산은 einsum('blhd,bmhd->blhm', Q, K) 및 einsum('blhm,bmhd->blhd', A, V)로 간결하게 구현된다. 배치 정규화의 통계량 계산, 합성곱의 축소 구조, 그리고 그래프 신경망의 메시지 전달 역시 모두 einsum 패턴으로 표현할 수 있다.
5. 최적화와 경로 선택
여러 텐서가 동시에 축소되는 복잡한 einsum 식에서는 축소 순서(contraction order)에 따라 전체 계산 비용이 크게 달라진다. 예를 들어 세 텐서 A_{ij}, B_{jk}, C_{k\ell}의 곱 \sum_{j,k} A_{ij} B_{jk} C_{k\ell}은 ((AB)C)의 순서로 계산하는 것과 (A(BC))의 순서로 계산하는 것의 비용이 일반적으로 다르다. 실제 구현에서는 opt_einsum과 같은 라이브러리가 탐욕적 또는 동적 계획법 기반의 경로 최적화 알고리즘을 사용하여 최적의 축소 순서를 탐색한다. 대규모 텐서 네트워크 계산에서는 이러한 최적화가 전체 실행 시간을 수 자릿수 이상 단축시킬 수 있다.
6. 통합적 표현의 의의
아인슈타인 합산 규약은 텐서 연산의 기저 수학적 구조를 명확하게 드러내며, 다양한 연산을 단일 문법으로 통합함으로써 코드의 가독성과 유지 보수성을 극대화한다. 또한 지표 수준의 표기는 명시적 형상 추적이 필요한 고차 텐서 연산의 정확성 검증을 용이하게 한다. 현대 딥러닝 프레임워크에서 einsum은 단순한 축약 도구를 넘어, 복잡한 텐서 조작을 선언적으로 기술하고 최적화된 커널 경로로 자동 매핑하는 통합적 표현 체계로 자리잡았다.