MNIST 손글씨 분류의 학습 흐름: 순전파·교차 엔트로피·역전파

인공지능 수학 공부에서 이미지 분류를 처음 만났을 때는 신경망, activation function, loss, gradient descent, backpropagation 같은 용어가 한꺼번에 등장했다. 각각을 따로 외우면 흐름이 끊긴다. MNIST 손글씨 분류를 기준으로 보면 모두 한 질문에 연결된다.

28×28 픽셀을 받은 model이 0부터 9까지의 class score를 만들고, 정답과의 차이를 줄이도록 parameter를 어떻게 바꾸는가?

이 글은 수학을 깊게 증명하기보다 한 번의 학습 step에서 tensor가 어떻게 이동하는지를 중심으로 정리한 학습 기록이다.

MNIST와 이미지 인식 단원을 공부하며 남긴 교재 사진

입력과 정답의 모양부터 잡는다

MNIST는 손글씨 숫자 이미지와 0~9 label로 구성된 dataset이다. torchvision.datasets.MNIST가 반환하는 원본 이미지는 PIL image이고, ToTensor()를 적용하면 보통 다음 모양이 된다.

한 이미지: [1, 28, 28]
batch 입력: [batch_size, 1, 28, 28]
label:      [batch_size]

완전연결층에 넣으려면 한 이미지를 784개 값으로 펼친다.

[batch_size, 1, 28, 28]
              ↓ flatten
[batch_size, 784]

모든 픽셀은 같은 의미를 갖지 않지만, 먼저 단순한 multilayer perceptron으로 전체 학습 흐름을 확인하기에는 충분하다.

신경망 layer가 계산하는 것

linear layer의 기본 계산은 다음과 같다.

z = xW + b

• x: 이전 layer가 만든 입력
• W: 학습할 weight
• b: 학습할 bias
• z: activation function을 적용하기 전 값

linear 계산만 여러 번 쌓으면 전체가 결국 하나의 linear transformation으로 합쳐질 수 있다. 중간에 ReLU 같은 nonlinear activation을 넣어야 더 복잡한 경계를 표현할 수 있다.

x
↓ Linear(784 → 128)
↓ ReLU
↓ Dropout
↓ Linear(128 → 10)
logits

ReLU는 max(0, x)를 계산한다. sigmoid처럼 큰 양수·음수 영역에서 gradient가 아주 작아지는 saturation 문제를 일부 줄이지만, 모든 vanishing-gradient 문제를 해결한다고 볼 수는 없다. 음수 영역의 출력과 gradient가 0이 되어 neuron이 계속 비활성화되는 문제도 생길 수 있다.

마지막 출력은 확률이 아니라 logits로 둔다

10개 숫자 class를 분류하므로 model의 마지막 출력은 sample마다 10개 score다.

logits shape = [batch_size, 10]

softmax를 적용하면 각 값을 0과 1 사이로 바꾸고 한 sample의 합을 1로 만들 수 있다.

softmax(z_i) = exp(z_i) / sum(exp(z_j))

다만 PyTorch의 CrossEntropyLoss를 사용할 때는 model 안에서 softmax를 먼저 적용하지 않는다. 공식 문서 기준으로 이 loss는 정규화되지 않은 logits를 입력으로 받으며, class index를 target으로 줄 때 log-softmax와 negative log-likelihood에 해당하는 계산을 안정적으로 묶어 처리한다.

logits = model(images)          # [batch_size, 10]
loss = criterion(logits, labels)  # labels: [batch_size]

예측 class만 필요하면 가장 큰 logit의 index를 고르면 된다. softmax는 확률 형태가 실제로 필요할 때 별도로 적용한다.

predicted = logits.argmax(dim=1)
probabilities = logits.softmax(dim=1)

교차 엔트로피가 줄어드는 방향을 찾는다

multi-class classification에서 cross entropy는 정답 class에 높은 score를 주도록 model을 학습시킨다. target이 class index인 한 sample의 직관적인 형태는 다음과 같다.

loss = -log(정답 class의 softmax 확률)

정답 class의 확률이 1에 가까워지면 loss는 0에 가까워지고, 0에 가까워지면 loss가 커진다.

MSE도 수학적으로 계산할 수 있지만, one-hot target과 sigmoid를 조합한 MSE가 MNIST 다중 분류의 기본 선택인 것은 아니다. 서로 배타적인 10개 class를 다룰 때는 logits와 cross entropy가 의도를 더 직접적으로 표현한다.

한 번의 학습 step

학습 loop의 핵심은 네 단계다.

1. optimizer에 남아 있는 gradient를 초기화
2. 순전파로 logits 계산
3. loss 계산 후 역전파
4. gradient를 이용해 parameter 갱신

optimizer.zero_grad()

logits = model(images)
loss = criterion(logits, labels)

loss.backward()
optimizer.step()

PyTorch는 순전파 중 tensor 연산의 계산 graph를 만들고, loss.backward()에서 chain rule을 따라 각 parameter의 gradient를 계산한다. optimizer는 그 gradient를 사용해 parameter를 갱신한다.

가장 단순한 stochastic gradient descent는 다음 방향으로 움직인다.

parameter := parameter - learning_rate × gradient

learning rate가 너무 크면 최솟값 근처를 지나치거나 발산할 수 있고, 너무 작으면 학습이 느리다. optimizer 이름만 바꾸기 전에 train loss와 validation metric이 어떻게 움직이는지 확인하는 것이 먼저다.

batch, iteration, epoch을 구분한다

• batch: 한 번의 순전파·역전파에 사용하는 sample 묶음
• iteration 또는 step: 한 batch로 parameter를 한 번 갱신한 것
• epoch: train dataset 전체를 한 번 사용한 것

sample이 60,000개이고 batch size가 128이라면 한 epoch은 대략 ceil(60000 / 128)개의 step으로 구성된다. 마지막 batch는 drop_last 설정에 따라 더 작거나 생략될 수 있다.

“SGD는 sample을 항상 한 개씩 사용한다”라고만 이해하면 framework의 optimizer 이름과 학습 방식이 뒤섞인다. PyTorch의 torch.optim.SGD도 mini-batch에서 평균·합산된 gradient로 흔히 사용한다.

Dropout과 train/eval mode

dropout은 학습 중 일부 activation을 무작위로 0으로 만들어 특정 unit에 지나치게 의존하지 않도록 돕는 regularization 방법이다. 학습과 평가의 동작이 다르기 때문에 mode 전환이 중요하다.

model.train()
train_logits = model(train_images)

model.eval()
with torch.no_grad():
    validation_logits = model(validation_images)

model.eval()은 dropout과 batch normalization처럼 train/eval behavior가 다른 module을 평가 mode로 바꾼다. gradient 계산 자체를 끄는 기능은 아니므로 inference에서는 torch.no_grad() 또는 torch.inference_mode()도 함께 사용한다.

실행 가능한 최소 예제

import torch
from torch import nn
from torch.utils.data import DataLoader, random_split
from torchvision import datasets
from torchvision.transforms import ToTensor

full_train_dataset = datasets.MNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)

test_dataset = datasets.MNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
)

train_dataset, validation_dataset = random_split(
    full_train_dataset,
    [55_000, 5_000],
    generator=torch.Generator().manual_seed(42),
)

train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
validation_loader = DataLoader(validation_dataset, batch_size=256)
test_loader = DataLoader(test_dataset, batch_size=256)

device = (
    "cuda" if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available()
    else "cpu"
)

model = nn.Sequential(
    nn.Flatten(),
    nn.Linear(28 * 28, 128),
    nn.ReLU(),
    nn.Dropout(p=0.2),
    nn.Linear(128, 10),
).to(device)

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)

for epoch in range(5):
    model.train()

    for images, labels in train_loader:
        images = images.to(device)
        labels = labels.to(device)

        optimizer.zero_grad()
        logits = model(images)
        loss = criterion(logits, labels)
        loss.backward()
        optimizer.step()

    model.eval()
    correct = 0
    total = 0

    with torch.no_grad():
        for images, labels in validation_loader:
            images = images.to(device)
            labels = labels.to(device)

            logits = model(images)
            predicted = logits.argmax(dim=1)

            correct += (predicted == labels).sum().item()
            total += labels.size(0)

    validation_accuracy = correct / total
    print(f"epoch={epoch + 1} validation_accuracy={validation_accuracy:.4f}")

model.eval()
test_correct = 0
test_total = 0

with torch.no_grad():
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)

        logits = model(images)
        predicted = logits.argmax(dim=1)
        test_correct += (predicted == labels).sum().item()
        test_total += labels.size(0)

print(f"final_test_accuracy={test_correct / test_total:.4f}")

이 코드는 train dataset 일부를 validation으로 분리해 epoch마다 확인하고, 학습이 끝난 뒤 test set을 한 번 평가하는 최소 예제다. accuracy 수치는 random seed, library version, device와 hyperparameter에 따라 달라지므로 특정 값을 보장하지 않는다.

normalization은 선택한 전처리다

ToTensor()는 image를 tensor로 바꾸고 일반적인 8-bit 입력을 [0, 1] 범위로 scale한다. Normalize(mean, std)는 channel마다 다음 변환을 적용한다.

output[channel] = (input[channel] - mean[channel]) / std[channel]

정규화가 언제나 학습을 빠르고 안정적으로 만든다고 단정할 수는 없다. 어떤 mean과 standard deviation을 어떤 dataset에서 계산했는지, train과 evaluation에 같은 transform을 적용했는지 기록해야 재현할 수 있다.

holdout과 cross-validation은 다르다

train/validation/test를 한 번 나누는 방식은 holdout split이다. k-fold cross-validation은 dataset을 k개 fold로 나누고 validation fold를 바꿔 가며 여러 번 학습·평가한다.

MNIST tutorial에서는 보통 정해진 train/test split을 사용하고, 필요하다면 train 일부를 validation으로 다시 나눈다. test set으로 hyperparameter를 반복 선택하면 최종 평가가 낙관적으로 치우칠 수 있다.

학습이 이상할 때 확인할 순서

1. image와 label shape이 예상과 같은가?
2. 마지막 layer가 10개의 logits를 내는가?
3. CrossEntropyLoss 앞에 softmax를 중복 적용하지 않았는가?
4. 매 step에서 zero_grad, backward, step 순서가 맞는가?
5. train에서는 model.train(), 평가에서는 model.eval()을 호출했는가?
6. train loss와 validation accuracy를 별도로 보고 있는가?
7. normalization과 random seed, version을 기록했는가?

선형 회귀에서 loss와 gradient의 관계를 먼저 보고 싶다면 인공지능 수학: 선형 회귀, 다른 입력 형태의 학습 문제는 자연어 처리 학습 기록과 이어진다.

처음에는 용어가 많아 보였지만 결국 반복되는 흐름은 단순하다. 입력을 score로 바꾸고, 정답과의 차이를 하나의 scalar loss로 만들고, 그 loss가 줄어드는 방향의 gradient로 parameter를 갱신한다. 각 기법이 이 흐름의 어느 지점에 놓이는지 확인하면 개념이 덜 흩어진다.

참고 자료

• PyTorch: MNIST dataset
• PyTorch: CrossEntropyLoss
• PyTorch: Softmax
• PyTorch: Module train and eval modes
• PyTorch: Normalize transform
• PyTorch quickstart tutorial