28.24 텐서 축소곱(Tensor Contraction)과 행렬 곱셈의 일반화

텐서 축소곱은 두 개의 텐서에서 특정 축(들)의 지표를 서로 일치시키고 해당 지표에 대해 합산을 수행함으로써 새로운 텐서를 생성하는 연산이다. 이 연산은 벡터의 내적, 행렬-벡터 곱, 행렬-행렬 곱, 그리고 고차 텐서 사이의 일반화된 곱셈을 하나의 통일된 수학적 틀 아래에서 기술할 수 있게 하며, 아인슈타인 합산 규약과 결합하여 매우 간결하고 강력한 표기법을 제공한다.

1. 축소곱의 형식적 정의

두 텐서 $A$ 와 $B$ 가 각각 차수 $p$ 와 $q$ 를 가지며, $A$ 의 축 $\alpha$ 와 $B$ 의 축 $\beta$ 의 크기가 동일하다고 하자. 이 두 축에 대한 축소곱은 다음과 같이 정의된다.

$C_{i_1 \cdots \hat{i}_\alpha \cdots i_p\, j_1 \cdots \hat{j}_\beta \cdots j_q} = \sum_{k=0}^{d-1} A_{i_1 \cdots i_{\alpha-1} k i_{\alpha+1} \cdots i_p}\, B_{j_1 \cdots j_{\beta-1} k j_{\beta+1} \cdots j_q}$

여기서 $\hat{i}_\alpha$ 와 $\hat{j}_\beta$ 는 해당 축이 축소되어 제외됨을 나타내며, $d$ 는 공통 축의 크기이다. 결과 텐서 $C$ 의 차수는 $p + q - 2$ 이다. 즉, 입력 텐서들의 차수 합에서 축소된 두 축을 뺀 값이다. 더 일반적으로 여러 쌍의 축을 동시에 축소할 수 있으며, $s$ 개의 축 쌍을 축소하면 결과의 차수는 $p + q - 2s$ 가 된다.

2. 특수 사례: 내적, 행렬 곱, 행렬-벡터 곱

텐서 축소곱의 단순한 특수 사례들을 살펴보면 그 일반화 원리가 명확해진다.

벡터 내적은 두 1차 텐서 $\mathbf{a}, \mathbf{b} \in \mathbb{R}^d$ 의 유일한 축을 축소하는 연산이다.

$\mathbf{a} \cdot \mathbf{b} = \sum_{k=0}^{d-1} a_k b_k$

결과는 0차 텐서(스칼라)이다.

행렬-벡터 곱은 행렬 $A \in \mathbb{R}^{M \times N}$ 의 두 번째 축과 벡터 $\mathbf{x} \in \mathbb{R}^N$ 의 유일한 축을 축소하는 연산이다.

$(A\mathbf{x})_i = \sum_{k=0}^{N-1} A_{ik} x_k$

결과는 1차 텐서(벡터)이다.

행렬-행렬 곱은 $A \in \mathbb{R}^{M \times K}$ 와 $B \in \mathbb{R}^{K \times N}$ 의 공통 차원 $K$ 를 축소한다.

$(AB)_{ij} = \sum_{k=0}^{K-1} A_{ik} B_{kj}$

결과 텐서의 형상은 $(M, N)$ 이다.

3. 고차 텐서 축소곱의 예시

형상 $(B, L, d)$ 인 3차 텐서 $A$ 와 형상 $(d, H)$ 인 2차 텐서 $W$ 의 축소곱은 $A$ 의 마지막 축과 $W$ 의 첫 번째 축을 축소하며, 결과는 다음과 같다.

$C_{b\ell h} = \sum_{k=0}^{d-1} A_{b \ell k} W_{k h}$

결과 텐서의 형상은 $(B, L, H)$ 이며, 이는 딥러닝의 완전 연결 계층이 배치 및 시퀀스 차원을 보존한 채 마지막 특성 차원만을 변환하는 구조에 정확히 대응한다. 여러 축을 동시에 축소하는 예로, 형상 $(B, C_{\text{in}}, H, W)$ 의 입력과 형상 $(C_{\text{out}}, C_{\text{in}}, K_h, K_w)$ 의 합성곱 커널 사이에서 $C_{\text{in}}$ 축을 축소하는 연산 등이 있다.

4. 아인슈타인 표기법과 축소곱

아인슈타인 합산 규약을 사용하면 축소곱을 매우 간결히 기술할 수 있다. 동일한 지표가 위와 아래(또는 동일 위치)에 반복되면 해당 지표에 대한 합산이 암묵적으로 수행된다. 위에서 제시한 예들은 각각 다음과 같이 표기된다.

연산	아인슈타인 표기	입력 차수	출력 차수
벡터 내적	$a_k b^k$	$1, 1$	$0$
행렬-벡터 곱	$A^i{}_k x^k$	$2, 1$	$1$
행렬-행렬 곱	$A^i{}_k B^k{}_j$	$2, 2$	$2$
트레이스	$A^i{}_i$	$2$	$0$
3차×2차 축소	$A_{b\ell k} W^k{}_h$	$3, 2$	$3$

5. 축소곱의 계산 복잡도

텐서 축소곱의 계산 비용은 축소된 축의 크기와 나머지 축의 크기의 곱에 비례한다. 예를 들어 $A \in \mathbb{R}^{M \times K}$ 와 $B \in \mathbb{R}^{K \times N}$ 의 축소곱은 $M \cdot N$ 개의 출력 원소 각각에 대해 $K$ 번의 곱셈과 합을 수행하므로 총 계산량은 $O(MNK)$ 이다. 일반적으로 형상 $(d_{i_1}, \ldots, d_{i_p})$ 과 $(d_{j_1}, \ldots, d_{j_q})$ 에서 $s$ 쌍의 축을 축소하는 연산의 복잡도는 다음과 같이 추정된다.

$\text{Cost} = O\!\left(\prod_{k \notin \text{contracted}} d_k \cdot \prod_{k \in \text{contracted}} d_k\right)$

여러 텐서가 동시에 축소되는 텐서 네트워크 계산에서는 축소 순서에 따라 총 비용이 크게 달라지므로, 최적 축소 순서 결정은 NP-난해 문제로 알려져 있다.

6. 행렬 곱셈의 일반화로서의 의미

텐서 축소곱은 선형대수의 행렬 곱셈을 다차원으로 완전히 일반화한 것으로서, 다중선형 사상(multilinear map)의 합성과 동일한 기하학적 의미를 갖는다. 이는 한 텐서가 정의하는 다중선형 사상의 출력 공간을 다른 텐서가 정의하는 사상의 입력 공간과 동일시함으로써, 두 사상을 합성하여 새로운 다중선형 사상을 얻는 연산에 해당한다. 딥러닝의 합성곱, 어텐션, 그리고 대부분의 선형 변환 계층은 본질적으로 텐서 축소곱의 특수한 형태로 표현될 수 있으며, 이러한 통일적 관점은 모델 구현과 최적화의 일관된 이해에 기여한다.