nn.Module — 작지만 진짜인 네트워크 · MLX Quest

처음부터 끝까지 한 번 돌려보자

이 레슨을 마치면 MLX로 완전한 학습 반복문 하나를 갖게 돼. 작은 2층 MLP를 정의하고, 직선에 잡음을 섞은 합성 회귀 자료를 만들고, SGD로 학습하면서 손실이 내려가는 걸 본 뒤 추론까지 해. 약 30줄이지만 장난감 흉내가 아니라 실제 학습 코드와 같은 모양이야.

목적은 흥미로운 모델 하나를 배우는 게 아니야. MLX 학습 반복문의 관절이 어떻게 맞물리는지 손으로 느끼는 거야. 뒤 트랙은 이 골격을 한 번 경험했다고 보고 진행해.

필요한 조각들

nn.Module — 모든 모델의 기반 클래스야. 상속한 뒤 하위 모듈을 만드는 __init__과 순전파를 맡는 __call__을 정의해. 별도 forward 메서드는 필요 없어. Python의 __call__이면 충분해.
nn.Linear, nn.relu — 이번에 쓸 층과 활성화 함수야.
mlx.optimizers.SGD — 가장 단순한 옵티마이저야. .update(model, grads)가 기울기를 모델에 제자리로 반영해.
nn.value_and_grad(model, loss_fn) — 한 번의 호출로 모델 매개변수에 대한 손실과 기울기를 함께 받는 표준 방식이야.
mx.eval(model.parameters(), optimizer.state) — 각 단계 끝에서 갱신된 가중치와 옵티마이저 상태를 실제 값으로 만들어. MLX가 한 단계를 묶어 최적화하도록 경계를 세워줘.

전체 반복문

아래 코드 블록이 학습 과정 전부야. 먼저 흐름이 얼마나 짧은지 읽고 직접 실행해. 손실값이 내려가는지 보고, 마지막 추론값이 실제 목표인 2.0 × 1.0 + 0.5 = 2.5에 제법 가까운지 확인해.

골격은 커져도 그대로야

작은 예제지만 모양은 진짜야. 트랙 5에서 70B 모델을 파인튜닝할 때도 모델 상속, nn.value_and_grad, 옵티마이저 갱신, 경계의 mx.eval을 그대로 써. 학습 대상이 커지면 반복문 주변의 복잡성은 늘어도 골격은 바뀌지 않아.

Code

처음부터 끝까지 — 합성 회귀 자료를 쓰는 2층 MLP·python

import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim


class MLP(nn.Module):
    def __init__(self, in_dim=1, hidden=16, out_dim=1):
        super().__init__()
        self.l1 = nn.Linear(in_dim, hidden)
        self.l2 = nn.Linear(hidden, out_dim)

    def __call__(self, x):
        return self.l2(nn.relu(self.l1(x)))


# Synthetic data: y = 2x + 0.5 + small noise
mx.random.seed(42)
N = 256
x = mx.random.uniform(-3, 3, (N, 1))
true_w, true_b = 2.0, 0.5
y = true_w * x + true_b + mx.random.normal((N, 1)) * 0.1

model = MLP()
mx.eval(model.parameters())   # materialize once before the loop

optimizer = optim.SGD(learning_rate=0.05)


def loss_fn(model, x, y):
    return ((model(x) - y) ** 2).mean()


loss_and_grad = nn.value_and_grad(model, loss_fn)

for step in range(200):
    loss_v, grads = loss_and_grad(model, x, y)
    optimizer.update(model, grads)
    mx.eval(model.parameters(), optimizer.state)
    if step % 50 == 0:
        print(f'step {step:3d}  loss {float(loss_v):.5f}')

# Verified output (2026-05-03):
#   step   0  loss 9.29228
#   step  50  loss 0.03690
#   step 100  loss 0.02353
#   step 150  loss 0.01696
#   final loss: 0.013964570127427578

# Quick inference check
xt = mx.array([[1.0]])
print('model(1.0) ≈', float(model(xt).item()), '  (true 2*1+0.5 = 2.5)')
# Verified: model(1.0) ≈ 2.544  — within noise of true target.

Exercise

학습 반복문을 실행해 손실이 내려가고 마지막 추론값이 2.5에 가까운지 확인해. 다음 두 변경을 한 번에 하나씩 시험해. (a) hidden을 64로 늘리고 다시 실행해. (b) mx.eval(model.parameters(), optimizer.state) 줄을 빼고 다시 실행해. (a)에서는 손실이 더 빨리 줄어? (b)에서도 학습이 돌아가고, Activity Monitor의 메모리는 어떻게 움직여? 두 문장으로 정리해. 두 번째 실험이 레슨 4의 지연 실행을 몸으로 보여줄 거야.

nn.Module — 작지만 진짜인 네트워크