[Lucid] 실제 텐서 연산들의 구현

🧮 실제 텐서 연산들의 구현

🔔 추상화가 코드 경로를 덜어낸 방식

지난 기록에서 @func_op가 forward/backward의 수학적 정의만 남기고 나머지 그래프 연결을 자동화하는 구조를 정리했다. 이번에는 그 틀을 기반으로 lucid/_func/bfunc.py, lucid/_func/ufunc.py에 있는 핵심 연산들이 어떻게 작성됐는지, gradient 수식과 코드 스니펫을 짝지어 기술한다. 목표는 딥러닝 루프에서 반복 호출되는 기본 연산을 일관된 템플릿 위에 올려놓는 것이었고, 실제 구현에서 고려한 broadcasting, shape 축소, CPU/GPU 분기 포인트를 함께 적는다.

🎛️ 기본 연산 선정 기준

연산을 추릴 때는 두 가지 조건을 걸었다.

1. 학습 루프에서 거의 매 스텝 호출되는가?
2. @func_op가 처리해야 할 그래프 후처리가 존재하는가?

이 조건을 만족하면서도 수식을 깔끔히 정리할 수 있는 다섯 가지를 골랐다: add, multiply, maximum, exp·log(쌍으로), matmul.

아래는 각 연산의 핵심 포인트와 구현 세부 내용이다.

➕ add: broadcasting과 gradient 정렬

덧셈은 간단하지만, 입력이 스칼라/텐서 혼합일 때 gradient shape을 맞춰 돌려주는지 확인해야 했다. 데코레이터가 입력 캐스팅과 그래프 연결을 처리하므로 수식만 명시한다.

# lucid/_func/bfunc.py
class add(operation):
    @binary_func_op()
    def cpu(self, a: Tensor, b: Tensor) -> _FuncOpReturnType:
        self.result = Tensor(a.data + b.data)
        return self.result, self.__grad__

    def __grad__(self) -> _GradFuncType:
        return self.result.grad, self.result.grad

• 수식: $y = a + b \Rightarrow \partial y/\partial a = 1,\ \partial y/\partial b = 1$
• 동작: 반환되는 gradient 두 개 모두 동일하며, _match_grad_shape가 내부에서 broadcasting된 축을 다시 접어 넣는다. 스칼라 덧셈, 행렬+벡터 등 다양한 조합을 같은 코드로 처리한다.

✖️ multiply: 입력 값이 곧 기울기

곱셈은 체인 룰의 기본 사례다. forward 시 입력을 캡처해 두고 backward에서 그대로 곱해준다. partial을 써서 입력을 안전하게 클로저에 넘겼다.

# lucid/_func/bfunc.py
class multiply(operation):
    @binary_func_op()
    def cpu(self, a: Tensor, b: Tensor) -> _FuncOpReturnType:
        self.result = Tensor(a.data * b.data)
        return self.result, partial(self.__grad__, a=a, b=b)

    def __grad__(self, a: Tensor, b: Tensor) -> _GradFuncType:
        return b.data * self.result.grad, a.data * self.result.grad

• 수식: $y = a \times b \Rightarrow \partial y/\partial a = b,\ \partial y/\partial b = a$
• 동작: gradient는 입력값과 결과 gradient를 곱한 형태로 전달된다. 입력이 broadcast됐다면 _match_grad_shape가 뒤에서 축을 압축해 leaf 텐서의 shape에 맞춘다.

🛡️ maximum: 마스크 기반 분기

maximum은 ReLU의 기반이 되므로 부호 분기 마스크가 정확해야 한다. forward에서 사용한 마스크를 backward에서도 동일하게 사용해 gradient를 분리한다.

# lucid/_func/bfunc.py
class maximum(operation):
    @binary_func_op()
    def cpu(self, a: Tensor, b: Tensor) -> _FuncOpReturnType:
        self.result = Tensor(np.maximum(a.data, b.data))
        return self.result, partial(self.__grad__, a=a, b=b)

    def __grad__(self, a: Tensor, b: Tensor) -> _GradFuncType:
        a_grad = (a.data >= b.data).astype(a.data.dtype)
        b_grad = (a.data < b.data).astype(b.data.dtype)
        return a_grad * self.result.grad, b_grad * self.result.grad

• 수식: $y = \max(a, b)$, $\partial y/\partial a = \mathbb{1}[a \ge b]$, $\partial y/\partial b = \mathbb{1}[a < b]$
• 동작: 부호 비교 결과를 float 마스크로 변환해 result.grad와 곱한다. ReLU(x) = \max(0, x)에선 $a=0$이 고정되므로 동일한 분기 로직이 그대로 적용된다.

🌡️ exp와 log: 수치 안정성의 기초

exp와 log는 softmax, cross-entropy의 수치 안정성에 직접 연결된다. unary 템플릿을 사용해 불필요한 인자 처리를 제거했고, 입력을 캡처해야 하는 경우만 partial을 사용했다.

# lucid/_func/ufunc.py
class exp(operation):
    @unary_func_op()
    def cpu(self, a: Tensor) -> _FuncOpReturnType:
        self.result = Tensor(np.exp(a.data))
        return self.result, self.__grad__

    def __grad__(self) -> _GradFuncType:
        return self.result.data * self.result.grad


class log(operation):
    @unary_func_op()
    def cpu(self, a: Tensor) -> _FuncOpReturnType:
        self.result = Tensor(np.log(a.data))
        return self.result, partial(self.__grad__, a=a)

    def __grad__(self, a: Tensor) -> _GradFuncType:
        return (1 / a.data) * self.result.grad

• 수식: $y = e^x \Rightarrow \partial y/\partial x = e^x$, $y = \log x \Rightarrow \partial y/\partial x = 1/x$
• 동작: exp는 forward 결과가 곧 미분 계수가 된다. log는 입력을 기억해야 하므로 partial을 통해 a를 클로저로 전달한다. 이 조합이 softmax에서 사용하는 input - max 안정화 경로에서 그대로 쓰인다.

🧭 matmul: 배치 차원과 축소 처리

행렬 곱은 입력 차원 정렬, 배치 broadcasting, 축소를 모두 포함한다. 동일한 수식을 CPU와 Metal GPU 모두에 적용하고, shape 축소는 helper로 분리했다.

# lucid/_func/bfunc.py
class matmul(operation):
    @binary_func_op()
    def cpu(self, a: Tensor, b: Tensor) -> _FuncOpReturnType:
        out = np.matmul(a.data, b.data)
        self.result = Tensor(out)
        return self.result, partial(self.__grad__, a=a, b=b, lib_=np)

    def __grad__(self, a: Tensor, b: Tensor, lib_: ModuleType) -> _GradFuncType:
        grad = self.result.grad
        if grad.ndim == 0:
            grad = lib_.reshape(grad, (1, 1))

        grad_a = lib_.matmul(grad, lib_.swapaxes(b.data, -1, -2))
        grad_b = lib_.matmul(lib_.swapaxes(a.data, -1, -2), grad)

        grad_a = self._reduce_broadcast_shape(grad_a, a.shape, lib_)
        grad_b = self._reduce_broadcast_shape(grad_b, b.shape, lib_)
        return grad_a, grad_b

• 수식 (2D): $Y = A B$, $\partial L/\partial A = (\partial L/\partial Y) B^\top$, $\partial L/\partial B = A^\top (\partial L/\partial Y)$
• 동작: 완전 축소된 경우 grad를 (1, 1)로 reshape해 연산을 통일한다. 배치 차원은 동일한 수식을 그대로 적용한 뒤 _reduce_broadcast_shape가 필요 없는 축을 sum(axis=...)로 제거해 원래 입력 shape에 맞춘다. backend 차이는 라이브러리 주입(np 또는 mx)으로만 구분한다.

🧵 정리

@func_op를 적용한 뒤 연산 구현은 수학적 정의와 필요한 캐시만 남고, 그래프 연결과 broadcasting 후처리는 템플릿이 처리하는 형태로 정리됐다. 덕분에

• gradient 수식이 잘못됐을 때만 수정하면 되고, 그래프 메타데이터는 건드리지 않는다.
• CPU/GPU 분기는 데코레이터 인자만 다르게 주면 되고, 코드 중복 없이 유지된다.
• _match_grad_shape와 _reduce_broadcast_shape가 공통 경로로 붙어 있어 shape 관련 버그를 조기에 차단한다.

다음 단계에서는 이 기반 위에서 컨볼루션, padding, 고차원 reduction 같은 연산을 확장하는 과정을 다룰 예정이다. 학습에 실제로 쓰이는 연산까지 동일한 템플릿으로 끌어올 수 있는지 점검하는 것이 목표다.