MATH / 중급

Σ와 첨자, 그리고 einsum: 수식을 그대로 코드로 옮기기

시그마와 첨자를 읽는 규칙을 자유 첨자·더미 첨자로 정리하고, 어텐션 점수와 합성곱 한 칸을 첨자로 완전히 풀어 쓴 뒤 einsum 한 줄로 되돌립니다.

PALDYN Team2026.08.0520 MIN READ

멀티헤드 어텐션의 점수를 배치까지 포함해 계산해야 합니다. 텐서는 (배치, 헤드, 토큰, 차원) 네 축이고, 마지막 축끼리 내적을 해서 (배치, 헤드, 토큰, 토큰)을 만들어야 합니다. 코드를 이렇게 씁니다.

scores = Q @ K.transpose(0, 1, 3, 2)

돌아갑니다. 그런데 잠시 뒤 헤드 축을 앞으로 빼는 리팩터링을 하고 나면 transpose(0, 1, 3, 2)가 맞는지 다시 확신이 안 섭니다. 축이 하나 늘거나 순서가 바뀔 때마다 이 자리를 다시 풀어야 하고, 그때마다 셰이프를 출력해 확인합니다. 코드는 도는데 왜 도는지를 매번 다시 확인하고 있는 상태입니다.

논문의 식에는 이 문제가 없습니다. 식은 처음부터 축의 순서를 안 씁니다. 대신 축마다 이름을 붙이고, 그 이름이 어디에 나타나는지로 계산을 정의합니다.

$(QK^\top)_{ij} = \sum_{t=1}^{d_k} Q_{it} K_{jt}$

지난 글에서 이 식의 기호들이 각각 무엇인지는 정리했습니다. 남은 것은 $\sum$ 과 첨자입니다. 어떤 첨자가 결과에 남고 어떤 첨자가 합쳐져 사라지는지 — 이 규칙 하나만 알면 위 식을 transpose 없이 그대로 코드로 옮길 수 있습니다.

Σ는 for 루프를 한 글자로 줄인 것이다

시그마의 정의는 간단합니다.

$\sum_{i=1}^{n} a_i = a_1 + a_2 + \cdots + a_n$

아래에 시작값, 위에 끝값, 오른쪽에 더할 것. 코드로 옮기면 그대로 루프입니다.

total = 0.0
for i in range(1, n + 1):
    total += a[i]

논문이 범위를 생략하는 일이 잦습니다. $\sum_i a_i$ 라고만 쓰면 " $i$ 가 가질 수 있는 값 전부"를 뜻하고, 어디까지인지는 $a$ 의 정의에서 옵니다. $a \in \mathbb{R}^n$ 이면 $1$ 부터 $n$ 까지입니다.

곱 버전도 있습니다. $\prod$ 는 더하는 대신 곱합니다.

$\prod_{i=1}^{n} p_i = p_1 p_2 \cdots p_n$

이 기호가 나오는 대표적인 자리가 언어 모델의 문장 확률입니다.

$p_\theta(x_1, \ldots, x_N) = \prod_{i=1}^{N} p_\theta(x_i \mid x_{<i})$

그리고 이 곱에 로그를 씌우면 합으로 바뀝니다 — $\log \prod = \sum \log$ . 사전학습 손실이 곱이 아니라 합으로 적혀 있는 이유가 그것이고, 왜 굳이 로그로 바꾸는지는 18번 글(지수와 로그)과 28번 글(최대가능도)이 답합니다. 지금은 곱을 합으로 바꾸는 통로가 로그라는 것만 기억하면 됩니다.

자유 첨자와 더미 첨자

첨자에는 두 종류가 있고, 이 구별이 이 글의 전부입니다.

$자유 첨자와 더미 첨자$

$(QK^\top)_{ij} = \sum_{t} Q_{it} K_{jt}$

$i$ 와 $j$ 는 자유 첨자입니다. 등호 왼쪽에도 나타나고 오른쪽에도 나타납니다. 결과에 남는 축이고, 그래서 결과 셰이프는 $(n, n)$ 입니다.
$t$ 는 더미 첨자입니다. $\sum$ 아래에만 있고 왼쪽에는 없습니다. 합쳐져 사라지는 축입니다.

여기서 두 가지 실용적인 규칙이 나옵니다.

규칙 1 — 더미 첨자는 이름을 바꿔도 값이 같습니다. $\sum_t Q_{it}K_{jt}$ 와 $\sum_z Q_{iz}K_{jz}$ 는 완전히 같은 식입니다. 루프 변수 이름을 바꾼 것과 다르지 않습니다. 논문마다 $t$ , $k$ , $d$ , $\ell$ 등을 제각각 쓰는 이유고, 이름이 다르다고 다른 식이라고 생각하면 안 됩니다.

규칙 2 — 자유 첨자는 등호 양변에 똑같이 나타나야 합니다. 왼쪽에 $i$ 가 있는데 오른쪽 어디에도 $i$ 가 없으면 그 식은 틀린 식입니다. 반대로 오른쪽에만 있고 왼쪽에 없는 첨자는 반드시 $\sum$ 에 묶여 있어야 합니다. 이것이 수식의 타입 검사입니다. 논문 식을 옮겨 적다가 첨자 하나를 빠뜨리면 이 규칙으로 즉시 잡힙니다.

규칙 2를 쓰면 결과 셰이프도 계산 없이 나옵니다. 왼쪽에 남은 자유 첨자의 개수가 곧 결과의 축 개수입니다. $(QK^\top)_{ij}$ 는 자유 첨자가 둘이니 2차원, $\sum_{ij} A_{ij}$ 는 자유 첨자가 없으니 스칼라입니다.

이중합과 순서 교환

합이 둘 겹치면 이렇게 씁니다.

$\sum_{i=1}^{m} \sum_{j=1}^{n} a_{ij}$

중첩 루프와 같습니다. 안쪽 $j$ 루프가 먼저 다 돌고, 그다음 바깥 $i$ 가 한 칸 움직입니다. 그런데 논문에서는 이 순서를 아무렇지 않게 바꿔 씁니다.

$\sum_{i}\sum_{j} a_{ij} = \sum_{j}\sum_{i} a_{ij}$

$이중합의 순서 교환$

왜 되는지는 그림이 답합니다. 두 식은 같은 칸들을 다른 순서로 훑을 뿐이고, 덧셈은 순서를 바꿔도 결과가 같습니다. 유한한 개수의 항을 더하는 한 언제나 성립합니다.

A = rng.normal(size=(9, 11))
A.sum(0).sum(0) == A.sum(1).sum(0)   # 같은 값

단, 유한할 때입니다. 항이 무한히 많으면 순서를 바꿔 값이 달라질 수 있고 — 조건수렴이라 부릅니다 — 그 경우를 구별하는 조건이 따로 있습니다. 딥러닝에서 다루는 합은 배치 크기, 토큰 수, 차원 수처럼 전부 유한하므로 마음 놓고 바꿔도 됩니다. 다만 적분이 섞이면 이야기가 달라져서, 기댓값과 미분의 순서를 바꾸는 자리(53번 글, 로그 미분 트릭)에서는 조건을 확인하고 넘어갑니다.

순서 교환이 실무에서 쓰이는 자리가 하나 있습니다. 배치 손실입니다.

$L = \frac{1}{BT}\sum_{b=1}^{B}\sum_{t=1}^{T} \ell_{bt}$

배치 안에서 먼저 평균 내고 토큰으로 평균 내든, 반대로 하든 같습니다. 다만 각 배치의 토큰 수가 다르면 이야기가 달라집니다. 패딩이 있는 배치에서 loss.mean(-1).mean()과 loss.sum() / mask.sum()이 다른 값을 주는 것이 그 때문입니다. 순서 교환이 깨진 게 아니라 나누는 수가 달라진 것인데, 첨자로 써 보면 어느 쪽이 의도한 식인지 바로 보입니다.

아인슈타인 합 규약

논문을 읽다 보면 $\sum$ 이 아예 없는 식을 만납니다.

$y_i = A_{ij} x_j$

오른쪽에 $j$ 가 두 번 나오는데 왼쪽에는 없습니다. 앞의 규칙 2에 따르면 이건 틀린 식이어야 합니다. 그런데 틀린 게 아니라 생략된 것입니다. 아인슈타인 합 규약은 이렇게 말합니다.

한 항 안에서 같은 첨자가 두 번 나타나면, 그 첨자에 대해 자동으로 합한다.

그래서 위 식은 $y_i = \sum_j A_{ij} x_j$ 와 같고, 이것은 행렬-벡터 곱입니다. 규약이 하는 일은 정확히 하나입니다 — 더미 첨자를 반복으로 표시하고 $\sum$ 을 지우는 것. 자유 첨자는 한 번만 나타나고 더미 첨자는 두 번 나타나므로, 세어 보면 어느 쪽인지 알 수 있습니다.

이 규약이 코드로 넘어온 것이 einsum입니다.

einsum 문자열 만드는 세 단계

np.einsum과 torch.einsum은 첨자 규약을 문자열로 받습니다. 만드는 절차는 기계적입니다.

입력 텐서마다 축에 글자를 붙인다. 왼쪽부터 순서대로, 축 하나에 글자 하나. 여러 텐서에서 같은 이름을 쓰면 그 축의 길이가 같다는 뜻입니다.
-> 오른쪽에 결과에 남길 글자를 적는다. 적은 순서가 곧 결과 축의 순서입니다.
오른쪽에 안 적은 글자는 합쳐진다.

3단계가 핵심입니다. -> 오른쪽에 안 적으면 사라지고, 사라진다는 것은 그 축에 대해 더한다는 뜻입니다. 자유 첨자와 더미 첨자의 구별이 문자열의 왼쪽·오른쪽으로 그대로 옮겨진 것입니다.

이 규칙 하나로 웬만한 텐서 연산이 다 표현됩니다.

연산	수식	einsum
내적	$\sum_i v_i w_i$	`np.einsum('i,i->', v, w)`
외적	$v_i w_j$	`np.einsum('i,j->ij', v, w)`
전치	$M_{ji}$	`np.einsum('ij->ji', M)`
행 합	$\sum_j M_{ij}$	`np.einsum('ij->i', M)`
대각 성분	$M_{ii}$	`np.einsum('ii->i', M)`
대각합	$\sum_i M_{ii}$	`np.einsum('ii->', M)`
행렬곱	$\sum_j A_{ij}B_{jk}$	`np.einsum('ij,jk->ik', A, B)`
원소별 곱	$A_{ij}B_{ij}$	`np.einsum('ij,ij->ij', A, B)`

표에서 대각 성분과 대각합을 비교해 보면 규칙이 선명해집니다. 'ii->i'는 반복된 글자를 오른쪽에 남겨 두었으니 합치지 않고 그 자리만 뽑고, 'ii->'는 안 남겼으니 합칩니다. 같은 왼쪽, 다른 오른쪽, 다른 연산 — 오른쪽이 무엇을 남길지 결정합니다.

손으로 따라가기 — QKᵀ를 첨자로 풀고 einsum으로 되돌리기

이제 처음의 어텐션 점수로 돌아갑니다. 먼저 첨자로 완전히 풀어 씁니다. 축 이름을 붙이는 것부터 시작합니다.

$Q$ 의 축: 토큰 $i$ , 차원 $t$ → it
$K$ 의 축: 토큰 $j$ , 차원 $t$ → jt
결과의 축: 질의 토큰 $i$ , 키 토큰 $j$ → ij

$K$ 의 토큰 축에 $i$ 가 아니라 $j$ 를 붙인 것이 유일한 판단입니다. 두 토큰 축은 길이가 같지만 같은 것을 세지 않습니다. $i$ 번 질의와 $j$ 번 키의 모든 조합이 필요하니 이름을 달리 줘야 하고, 만약 둘 다 $i$ 로 두면 아인슈타인 규약이 그 축을 합쳐 버려 대각 성분만 남게 됩니다.

차원 축은 둘 다 $t$ 입니다. 같은 이름을 두 번 썼고 오른쪽에 안 적었으니 합쳐집니다. 그것이 내적입니다.

$(QK^\top)_{ij} = \sum_{t} Q_{it}K_{jt} \quad\longrightarrow\quad \texttt{'it,jt->ij'}$

세 겹 루프로 직접 짠 것과 같은지 확인합니다.

import numpy as np

rng = np.random.default_rng(0)
n, dk = 4, 6
Q = rng.normal(size=(n, dk))
K = rng.normal(size=(n, dk))

S = np.zeros((n, n))
for i in range(n):
    for j in range(n):
        total = 0.0
        for t in range(dk):
            total += Q[i, t] * K[j, t]
        S[i, j] = total

print(np.allclose(S, np.einsum('it,jt->ij', Q, K)))   # True
print(np.allclose(S, Q @ K.T))                        # True
print(np.allclose(np.einsum('it,jt->ij', Q, K),
                  np.einsum('iz,jz->ij', Q, K)))      # True — 더미는 이름 무관

세 줄 모두 True입니다. 마지막 줄이 규칙 1의 확인입니다.

이제 배치와 헤드를 붙입니다. 실제 텐서는 (B, H, T, D) 네 축입니다. 첨자를 앞에 두 개 더 붙이기만 하면 됩니다.

$S_{bhij} = \sum_{d} Q_{bhid} K_{bhjd} \quad\longrightarrow\quad \texttt{'bhid,bhjd->bhij'}$

B, H, T, D = 2, 3, 5, 7
Qb = rng.normal(size=(B, H, T, D))
Kb = rng.normal(size=(B, H, T, D))

Sb = np.einsum('bhid,bhjd->bhij', Qb, Kb)
print(Sb.shape)                                        # (2, 3, 5, 5)
print(np.allclose(Sb, Qb @ Kb.transpose(0, 1, 3, 2)))  # True

두 줄이 같은 값을 줍니다. 다른 것은 읽는 데 드는 노력입니다. transpose(0, 1, 3, 2)는 축 번호를 세어야 검증되지만, 'bhid,bhjd->bhij'는 논문 식의 첨자를 그대로 옮겨 적은 것이라 눈으로 대조하면 끝납니다. 배치 축이 하나 더 늘어도 글자 하나를 앞에 붙이면 되고, 헤드 축을 뒤로 옮겨도 문자열 안의 글자 위치만 바꾸면 됩니다. 축 순서를 다시 계산할 일이 없어집니다.

합성곱 한 칸도 같은 방법으로

같은 절차가 합성곱에도 그대로 적용됩니다. 출력 한 칸의 정의는 이렇습니다.

$Y_{n,o,y,x} = \sum_{c}\sum_{k_y}\sum_{k_x} X_{n,\,c,\,y+k_y,\,x+k_x}\; W_{o,\,c,\,k_y,\,k_x}$

자유 첨자는 $n$ (배치), $o$ (출력 채널), $y$ , $x$ (출력 위치) 넷이고, 더미 첨자는 $c$ (입력 채널), $k_y$ , $k_x$ (커널 안의 위치) 셋입니다. 결과가 4차원인 이유가 자유 첨자 넷이고, 세 겹 합이 붙은 이유가 더미 첨자 셋입니다.

한 가지 걸리는 것이 $X$ 의 첨자에 있는 $y+k_y$ 입니다. 첨자 자리에 덧셈이 들어가 있어서 einsum 문자열로 바로 못 옮깁니다. 이럴 때는 겹치는 창을 먼저 꺼내 축으로 만들어 둡니다.

N, C, Hh, Ww = 2, 3, 6, 6
O, KH, KW = 4, 3, 3
X = rng.normal(size=(N, C, Hh, Ww))
W = rng.normal(size=(O, C, KH, KW))

# (N, C, OH, OW, KH, KW) — y+ky 를 미리 펼쳐 축으로 만든다
patches = np.lib.stride_tricks.sliding_window_view(X, (KH, KW), axis=(2, 3))
print(patches.shape)                       # (2, 3, 4, 4, 3, 3)

Y = np.einsum('ncyxkl,ockl->noyx', patches, W)
print(Y.shape)                             # (2, 4, 4, 4)

문자열을 읽어 봅니다. ncyxkl에서 n, y, x는 오른쪽에 남았고 c, k, l은 안 남았으니 합쳐집니다. ockl에서 o는 남고 c, k, l은 두 입력에 모두 나타나며 안 남았으니 합쳐집니다. 수식의 자유 첨자 넷과 더미 첨자 셋이 문자열에 그대로 있습니다.

이것이 여섯 겹 루프와 같은지 확인하면 True가 나옵니다. 그리고 이 형태가 실제로 합성곱을 행렬곱으로 구현하는 방법 — im2col — 의 뼈대입니다. 첨자를 풀어 쓰는 순간 "합성곱은 사실 큰 행렬곱이다"라는 문장이 계산으로 보입니다.

einsum을 쓸 때 알아 둘 것

축 순서를 바꾸는 데는 공짜가 아닙니다. 'bhid,bhjd->bijh'처럼 결과 축 순서를 바꾸면 메모리 재배치가 일어납니다. 문자열이 짧다고 연산이 싼 것은 아닙니다.

세 개 이상을 한 번에 곱할 때는 순서가 성능을 가릅니다. np.einsum('ij,jk,kl->il', A, B, C)는 $(AB)C$ 로 계산하느냐 $A(BC)$ 로 계산하느냐에 따라 연산량이 크게 달라집니다. NumPy는 기본적으로 왼쪽부터 순서대로 하므로 optimize=True를 주어야 최적 순서를 찾습니다. PyTorch의 torch.einsum은 기본으로 경로 최적화를 합니다. 왜 결합 순서가 연산량을 바꾸는지는 10번 글(행렬곱과 셰이프 산술)에서 계산으로 확인합니다.

반복된 첨자를 오른쪽에 남기는 것은 합이 아니라 추출입니다. 'ii->i'가 대각 성분을 뽑는다는 것을 앞에서 봤습니다. 이 동작은 아인슈타인 규약 원래 정의에는 없고 einsum이 확장한 부분이라, 논문 식을 옮길 때는 "이 반복이 합인가 추출인가"를 한 번 확인하는 게 좋습니다.

다시 그 코드로

처음의 코드로 돌아갑니다.

scores = Q @ K.transpose(0, 1, 3, 2)

이 줄이 왜 맞는지 이제 축 번호를 세지 않고 답할 수 있습니다. 그리고 애초에 이렇게 쓰면 다시 확인할 일이 없습니다.

scores = np.einsum('bhid,bhjd->bhij', Q, K)

논문의 $\sum_d Q_{bhid}K_{bhjd}$ 와 이 문자열은 같은 것을 다른 문법으로 쓴 것입니다. 축의 순서가 아니라 축의 이름으로 계산을 지정하기 때문에, 텐서 레이아웃이 바뀌어도 식은 그대로입니다.

여기까지가 수식을 읽는 데 필요한 문법입니다. 글자 모양으로 자료형을 읽고, 첨자로 축을 읽고, 시그마로 무엇이 사라지는지 읽습니다. 이제 문법 말고 대상 차례입니다. 임베딩 한 줄이 왜 그냥 float 배열이 아니라 벡터인지 — 그 구별이 무엇을 보장하길래 "king − man + woman"이 말이 되는지 — 가 다음 글의 질문입니다.

지난 글: 논문 수식 표기 규약: 스칼라·벡터·행렬·텐서를 구별해서 읽기

다음 글: 벡터: 임베딩 한 줄이 숫자 배열이 아니라 벡터인 이유

읽어주셔서 감사합니다. 😊

목록으로 전체 글

KEEP READING