[Lucid] Functional API 구축

🧱 Functional API 구축

🧭 계층적 추상화 설계

Lucid를 설계할 때 가장 먼저 박은 말뚝은 계층별로 책임을 분리한다는 원칙이었다. 아래와 같은 흐름을 엄격하게 유지하려 했다.

1. Tensor: 데이터/gradient를 들고 있는 최하위 노드.
2. 기본 연산: add, mul, matmul 등 수학적 primitive. 여기서만 NumPy가 직접 호출된다.
3. nn.functional: 선형 변환, 활성화, 풀링, 정규화, 드롭아웃, 손실 등 상태 없는 연산을 primitive 조합으로 정의한다.
4. nn.Module: Conv2d, Linear, BatchNorm 등 상태(state)를 가진 모듈.
5. 모델: LeNet, AlexNet 등 실제 네트워크를 모듈 조합으로 정의.

이 계층을 지키면 유지보수성과 시스템적 가독성이 확보된다. 특히 NumPy 호출이 2단계(기본 연산)에서만 발생하도록 봉인하면, 상위 레이어는 모두 동일한 추상화 규칙에 의존하게 된다. functional 계층을 만들면서도 이 규율을 무조건 지키는 것을 목표로 삼았다. 기능별 API를 작성할 때마다 “이 코드가 primitive 외부 호출을 섞지 않는가?”, “모듈이 가져가야 할 상태를 functional에 숨기지 않았는가?” 같은 체크리스트를 반복했다.

---

🧩 nn.functional의 목표와 제약

nn.functional은 상태 없는 순수 연산 집합이다. 입력 Tensor와 파라미터 Tensor를 받아 즉시 출력 Tensor를 반환하며, 상태 저장을 하지 않는다. 따라서

• autograd: gradient 추적은 이미 하위 연산(@func_op 기반)에서 처리되므로 별도 로직이 필요 없다.
• 디바이스 종속성 최소화: NumPy 외부 호출 없이, 상위 API에서는 오직 하위 primitive(lucid.add, lucid.matmul, lucid.exp 등)만 사용한다.
• 입출력 규약 명확화: broadcasting, shape 변환, 패딩/스트라이드 계산을 함수 단위로 캡슐화해 모듈(nn.Conv2d 등)이 그대로 재사용 가능하도록 한다.

이 목표를 지키기 위해 대표 functional 연산들을 정리한다. relu, softmax, linear, max_pool, batch_norm, dropout, cross_entropy처럼 학습 루프에서 반복 호출되는 API가 하위 primitive만으로 어떻게 구성되는지가 포인트다. (컨볼루션은 구현이 길어 별도 문서에서 다룰 예정이다.) 각 함수가 primitive 호출만으로 작성되면, 바로 그 위 층(nn.Module)은 상태 보관 + functional 위임이라는 얇은 역할에 집중할 수 있다.

⚡ relu: 마스크 기반 분기

경로: lucid/nn/functional/_activation.py

def relu(input_: Tensor) -> Tensor:
    return lucid.maximum(0, input_)

• 의존성: maximum(primitive)만 사용. NumPy 호출 없음.
• 수식: $f(x)=\max(0, x)$, $\partial f/\partial x = \mathbb{1}[x>0]$
• 설계 메모: 분기 마스크를 primitive가 처리하므로 functional에서는 로직 분기가 없다. 상위 nn.ReLU 모듈은 이 함수를 그대로 호출해 forward를 구성한다.

🌡️ softmax: 안정화와 축 규약

경로: lucid/nn/functional/_activation.py

def softmax(input_: Tensor, axis: int = -1) -> Tensor:
    input_max = lucid.max(input_, axis=axis, keepdims=True)
    input_stable = input_ - input_max
    e_input = lucid.exp(input_stable)
    sum_e_input = e_input.sum(axis=axis, keepdims=True)
    return e_input / sum_e_input

• 의존성: max, sub, exp, sum, truediv 등 모두 primitive.
• 수식: $\mathrm{softmax}(x)_i = \exp(x_i - \max x) / \sum_j \exp(x_j - \max x)$
• 설계 메모: 수치 안정화를 함수 내부에서 처리해, 상위 모듈이나 loss 구현에서 별도 전처리를 하지 않아도 된다. gradient는 하위 연산들이 조합돼 자동 계산된다.

🧮 linear: 행렬 곱과 편향

경로: lucid/nn/functional/_linear.py

def linear(input_: Tensor, weight: Tensor, bias: Tensor | None = None) -> Tensor:
    output = input_ @ weight.mT
    if bias is not None:
        output += bias
    return output

• 의존성: matmul(primitive)과 add.
• 수식: $Y = X W^\top + b$
• 설계 메모: 편향은 broadcasting으로 더해지므로 별도 reshape 없이 primitive 규약에 맡긴다. 이 함수가 nn.Linear.forward의 전부이며, 파라미터는 모듈이 소유하고 연산은 functional이 담당하는 구조가 유지된다.

🧪 batch_norm: 러닝 스탯과 브로드캐스트

경로: lucid/nn/functional/_norm.py

def batch_norm(input_, running_mean, running_var, weight=None, bias=None, training=True, momentum=0.1, eps=1e-5):
    use_batch_stats = training or running_mean is None or running_var is None
    if use_batch_stats:
        batch_mean = input_.mean(axis=(0, *range(2, input_.ndim)), keepdims=True)
        batch_var = input_.var(axis=(0, *range(2, input_.ndim)), keepdims=True)
        mean, var = batch_mean, batch_var
    else:
        mean = running_mean.reshape(1, C, *(1,) * spatial_dim)
        var = running_var.reshape(1, C, *(1,) * spatial_dim)

    normalized = (input_ - mean) / lucid.sqrt(var + eps)
    if weight is not None:
        normalized *= weight.reshape((1, C) + (1,) * spatial_dim)
    if bias is not None:
        normalized += bias.reshape((1, C) + (1,) * spatial_dim)
    return normalized

• 의존성: mean, var, sqrt, sub, div, mul, add 등 primitive만 사용.
• 수식: $\hat{x} = (x - \mu) / \sqrt{\sigma^2 + \epsilon}$, $y = \gamma \hat{x} + \beta$
• 설계 메모: 러닝 스탯 업데이트는 모듈이 소유하지만, 연산 자체는 primitive 조합으로 끝낸다. spatial 차원 수를 자동으로 맞추기 위해 reshape 패턴을 일관되게 적용했다.

🌧️ dropout: 마스킹과 스케일링

경로: lucid/nn/functional/_drop.py

def dropout(input_: Tensor, p: float = 0.5, training: bool = True) -> Tensor:
    if not training:
        return input_
    mask = (lucid.random.rand(*input_.shape) > p).free()
    scale = 1.0 / (1 - p)
    return input_ * mask * scale

• 의존성: random.rand, mul, sub 등 primitive만 사용.
• 수식: $\mathrm{dropout}(x) = \dfrac{\mathbb{1}[u>p]}{1-p} \cdot x,\ u\sim\mathcal{U}(0,1)$
• 설계 메모: 학습/평가 모드를 분리하고, mask를 .free()로 그래프에서 떼어 gradient가 흘러가지 않게 했다. spatial 차원을 단일 축으로 묶는 dropoutnd도 동일한 패턴으로 구현했다.

🏊 max_pool: 슬라이스 재사용

경로: lucid/nn/functional/_pool.py

def max_pool2d(input_, kernel_size=1, stride=1, padding=0) -> Tensor:
    return lucid.max(_pool2d(input_, kernel_size, stride, padding), axis=-1)

• 의존성: 내부 _pool2d는 pad, stack, slice 조합만 사용. 최종 축소는 primitive max.
• 수식: 각 윈도우에서 $\max$ 선택.
• 설계 메모: 1D/2D/3D를 동일 패턴으로 구현하고, 커널/스트라이드/패딩을 튜플로 정규화해 모듈(nn.MaxPool2d)이 그대로 위임할 수 있다.

🧭 layer_norm: 축 기준 정규화

경로: lucid/nn/functional/_norm.py

def layer_norm(input_, normalized_shape, weight=None, bias=None, eps=1e-5):
    mean = input_.mean(axis=tuple(range(-len(normalized_shape), 0)), keepdims=True)
    var = input_.var(axis=tuple(range(-len(normalized_shape), 0)), keepdims=True)
    normalized = (input_ - mean) / lucid.sqrt(var + eps)
    if weight is not None:
        normalized *= weight.reshape((1,) * (input_.ndim - len(normalized_shape)) + normalized_shape)
    if bias is not None:
        normalized += bias.reshape((1,) * (input_.ndim - len(normalized_shape)) + normalized_shape)
    return normalized

• 의존성: mean, var, sqrt, sub, div, mul, add.
• 수식: 입력의 마지막 normalized_shape 축들에 대해 평균/분산을 구해 정규화 후 affine.
• 설계 메모: 배치 크기나 시퀀스 길이에 상관없이 동일한 패턴을 적용하기 위해 reshape 패턴을 일반화했다.

🎯 cross_entropy: 안정화된 확률 변환

경로: lucid/nn/functional/_loss.py

def cross_entropy(input_, target, weight=None, reduction="mean", eps=1e-7, ignore_index=None):
    exp_logits = lucid.exp(input_ - lucid.max(input_, axis=1, keepdims=True))
    prob = exp_logits / lucid.sum(exp_logits, axis=1, keepdims=True)
    loss = -lucid.log(prob[indices, target_int] + eps)
    ...
    return _loss_reduction_or_ignore(loss, weight, target_int, ignore_index, reduction)

• 의존성: max, exp, sum, truediv, log, 인덱싱 등 primitive만 사용.
• 수식: $\mathrm{CE}(z, y) = -\log \dfrac{\exp(z_y)}{\sum_j \exp(z_j)}$ (안정화를 위해 $z - \max z$ 적용)
• 설계 메모: 안정화(shift)와 마스킹(ignore_index)을 함수 내부에서 처리해 상위 모듈이나 학습 루프가 단순화된다.

---

🧵 마무리 및 다음 단계

nn.functional을 구성하면서 primitive 레이어에 모든 저수준 의존성을 몰아넣는 규율이 작동하는지 검증했다. 결과적으로:

• 상위 API는 오직 Lucid primitive 호출만으로 작성되어 코드 경로가 읽기 쉽고 테스트 범위도 좁혀졌다.
• NumPy 직접 호출이 하위 연산에만 존재하므로, backend를 교체하거나 확장할 때 수정 범위가 명확해진다.
• 모듈(nn.Conv2d, nn.Linear, nn.LayerNorm, nn.Dropout, nn.MaxPool2d 등)은 상태만 보유하고 forward는 functional 호출로 구성되어 역할 분리가 유지된다.
• 테스트 단위가 명확해져 primitive → functional → module → model 순으로 점진적 검증이 가능하다.

다음 기록에서는 (별도 문서로 분리하는) 컨볼루션 구현을 먼저 다루고, 이어서 이 functional 레이어 위에 얹힌 nn.Module 구현과, 이를 합쳐 간단한 모델(예: LeNet)까지 빌드하는 흐름을 정리할 예정이다. 이렇게 계층적으로 쌓은 추상화가 실제 학습 파이프라인에서도 일관성을 유지하는지 점검하는 것이 목표다.