수학: 원소별, 행렬곱, 브로드캐스팅

텐서 수학의 세 범주

원소별 연산

표준 연산자 + - * / **와 exp, log, sqrt, abs, clamp 같은 대부분의 torch.foo 함수는 원소별로 동작해. a * b의 별표는 행렬 곱셈이 아니라 원소별 곱셈이야. 교과서의 수식을 코드로 옮길 때 가장 자주 헷갈리는 지점이지.

행렬 곱셈

행렬 곱셈에는 Python 3.5에서 PEP 465로 도입된 @ 연산자나 같은 뜻의 torch.matmul을 써. 둘 다 2D 행렬끼리의 곱과 배치로 묶인 3D 이상 텐서의 곱을 지원하고, 앞쪽 차원은 필요에 따라 브로드캐스트해. torch.bmm은 정확히 3D 텐서만 받으므로 모양 계약을 엄격하게 지키고 싶을 때 유용해.

축소 연산

sum, mean, max, argmax, std 같은 축소 연산은 하나 이상의 차원을 줄여. 어느 차원을 줄일지는 dim 인자가 정해. x.sum()은 모든 원소를 더하고, x.sum(dim=0)은 0번 차원을 따라 더해. keepdim의 기본값은 False라 축소한 차원이 사라져. 뒤 연산에서 브로드캐스팅하기 쉽게 유지하려면 keepdim=True를 써.

브로드캐스팅

브로드캐스팅은 크기가 1인 차원을 가상으로 늘려 서로 다른 모양의 텐서를 결합하게 해 줘. 오른쪽 차원부터 다음 규칙을 확인해:

서로 대응하는 차원의 크기가 같거나, 둘 중 하나가 1이어야 해.
한쪽에 없는 앞쪽 차원은 크기가 1인 것으로 취급해.

NumPy와 같은 규칙이야. 헷갈리면 두 모양을 오른쪽에 맞춰 위아래로 적고, 오른쪽부터 한 쌍씩 확인해.

Code

원소별 대 행렬 곱하기·python

import torch

a = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
b = torch.tensor([[5.0, 6.0], [7.0, 8.0]])

# ELEMENT-WISE — what the asterisk does
a * b
# tensor([[ 5., 12.],
#         [21., 32.]])

# MATRIX MULTIPLY — what @ does
a @ b
# tensor([[19., 22.],
#         [43., 50.]])

# Both also exist as named functions
torch.mul(a, b)        # element-wise
torch.matmul(a, b)     # matrix multiply

배치로 묶은 행렬곱 (어텐션 방식)·python

import torch

# Q, K, V in attention: (batch, heads, seq, head_dim)
Q = torch.randn(2, 8, 64, 32)
K = torch.randn(2, 8, 64, 32)

# Attention scores: (batch, heads, seq, seq)
# K.transpose(-2, -1) → (2, 8, 32, 64)
scores = (Q @ K.transpose(-2, -1)) / (32 ** 0.5)
print(scores.shape)   # torch.Size([2, 8, 64, 64])

# torch.bmm is the strictly-3D version (no broadcasting on the leading dim)
A = torch.randn(8, 3, 4)
B = torch.randn(8, 4, 5)
torch.bmm(A, B).shape  # torch.Size([8, 3, 5])

축소와 keepdim·python

import torch

t = torch.tensor([[1.0, 2.0, 3.0],
                  [4.0, 5.0, 6.0]])

t.sum()             # tensor(21.) — over everything
t.sum(dim=0)        # tensor([5., 7., 9.]) — collapse rows → shape (3,)
t.sum(dim=1)        # tensor([6., 15.])     — collapse cols → shape (2,)

# keepdim=True preserves the dim, which keeps broadcasting valid downstream
mean_per_row = t.mean(dim=1, keepdim=True)   # shape (2, 1)
centered = t - mean_per_row                  # broadcasts cleanly
print(centered)

실전 브로드캐스팅·python

import torch

t = torch.zeros(3, 4)
row = torch.tensor([1, 2, 3, 4])      # (4,)        broadcasts down rows
col = torch.tensor([[10], [20], [30]])  # (3, 1)    broadcasts across cols

(t + row).shape   # torch.Size([3, 4])
(t + col).shape   # torch.Size([3, 4])
(t + row + col).shape  # torch.Size([3, 4])

# Right-aligned shape check:
#         (3, 4)
#            (4,)   ← matches col 4, missing dim treated as 1
#         (3, 1)    ← matches col 4 via 1, row matches 3

Exercise

(batch=2, heads=4, seq=8, head_dim=16) 모양의 Q, K, V 텐서로 어텐션을 직접 구현해. 점수를 계산해 sqrt(head_dim)으로 나누고, 마지막 차원에 소프트맥스를 적용한 뒤 어텐션이 반영된 값을 구해. 출력 모양을 현대적인 PyTorch의 torch.nn.functional.scaled_dot_product_attention과 비교해. 정확히 같아야 해.