[Lucid] Tensor 클래스의 구현

🧱 Tensor 클래스의 구현

🧑‍💻 추상적 아이디어에서 실제 코드로

앞선 기록에서 나는 Lucid 프로젝트의 출발점이 되었던 autodiff의 추상적 구조인 Tensor를 node로 바라보고 computation graph를 구성한 뒤, 이를 역순으로 순회하며 gradient를 전파한다는 아이디어를 정리했다. 이제 그 다음 단계에서 할 일은 그 개념적 설계를 실제 코드 구조로 옮기는 것이다.

Lucid의 핵심 아이디어는 어디까지나 Tensor 클래스이며, 이는 lucid._tensor.tensor.Tensor namespace 안의 구현을 통해 구체적으로 드러냈다. Tensor는 단순히 배열을 담는 그릇이 아니라, 계산 그래프의 노드이머 gradient의 출발점이자 종착점이다. 따라서 이 클래스가 어떻게 구성되는지가 Lucid 전체의 철학과 직결된다. (편의상 lucid.Tensor를 통해 최상위 namespace에서 접근할 수 있게 만들었다.)

이 글에서는 Lucid의 Tensor가 어떤 방식으로 autodiff 구조를 반영하며, 특히 backward 메커니즘이 어떤 형태로 다듬어졌는지를 중점적으로 기술할 예정이다. 모든 메서드와 속성을 나열하는 식의 문서가 아니라, 개발 과정에서 내가 어떤 문제를 해결하고자 했는지, 어떤 결정이 Tensor의 핵심 구조에 반영되었는지 를 담아내려고 한다.

---

🏗️ Tensor 클래스의 기본 구조를 설계하며

추상적으로 구상한 Tensor를 실제 코드로 구현하는 첫 단계에서 나는 "어떤 signature를 가진 객체여야 하는가" 부터 정리했다. Lucid의 초기 Tensor 생성자는 대략 다음과 같은 형태를 하고 있다.

class Tensor:
	def __init__(
    	self, data: _ArrayOrScalar, 
        requires_grad: bool = False, keep_grad: bool = False
	) -> None:
    	# 실제 값 저장
    	self.data = np.array(data)
        
        # computation graph 관련 메타데이터
        self._prev: list[Tensor] = []
        self._op: Optional[_OpBase] = None
        
        # autodiff 관련 상태
        self.grad: Optional[np.ndarray] = None
        self.requires_grad = requires_grad
        self.keep_grad = keep_grad

여기서 중요한 포인트는 data 자체보다도, 그 주변을 감싸고 있는 autodiff 메타데이터 층(layer)이다. self._prev와 self._op는 computation graph를 구성하는 핵심 축이고, grad, requires_grad, keep_grad는 gradient 추적과 해제의 책임을 맡는다.

Tensor를 설계할 때 가장 먼저 고려했던 점은 "데이터 + 그래프 정보 + gradient"라는 세 가지 요소를 어떻게 조화롭게 묶어낼 것인가였다. Lucid Tensor는 PyTorch Tensor와 비슷하게 사용되지만, 내부적으로는 NumPy의 ndarray를 기반으로 한다. 이 ndarray는 텐서의 실제 값을 나타내는 본질적 데이터이다.

하지만 Tensor는 단순한 데이터 컨테이너가 아니라, 다음과 같은 추가적인 메타 정보 를 반드시 포함해야 했다.

• _prev: 자신을 생성한 부모(parent) 텐서들
• _op: 자신을 생성한 operation 객체
• grad: 역전파 시 누적될 gradient
• requires_grad: 이 텐서가 gradient 추적 대상인지 여부
• is_leaf: 파라미터나 입력처럼 그래프의 말단에 위치하는지 여부

이 중에서 개발 과정에서 특별히 고민을 많이 했던 요소는 _prev, _op, 그리고 requires_grad였다.

🔗 부모 텐서와 연산의 구조화

앞서 구상했던 것처럼, Tensor가 하나의 node가 되기 위해서는 연산 과정에 대한 인과관계를 명확히 기록해야 한다. 따라서 Tensor를 생성하는 거의 모든 operation은 다음과 같은 패턴을 따른다.

1. 연산이 수행되며 새로운 Tensor 결과값을 생성한다.
2. 생성된 Tensor는 _prev에 해당 연산의 입력 텐서들을 기록한다.
3. _op에는 해당 연산을 담당한 operation 객체를 기록한다.

마지막의 두 가지 정보는 backward 시 매우 중요한 역할 을 한다. 역전파는 결국 자신을 생성한 연산의 local derivative를 부모에게 적용하는 과정이므로, _op는 backward의 수학적 핵심이고, _prev는 그래프 traversal의 핵심이라 할 수 있다.

처음에는 Tensor가 단순히 부모 리스트만 가지고 있어도 충분하다고 생각했지만, 실제 backward 메커니즘을 구성하기 위해서는 연산이 스스로 gradient rule을 알고 있어야 했다. 그래서 모든 operation을 하나의 클래스로 추상화하여, 각 연산 객체가 자신만의 backward() 메서드를 가지도록 설계하였다.

이 구조는 "Tensor는 node, operation은 edge의 성격을 띤다" 라는 비유적인 기반에서 출발했다.

🔍 requires_grad 속성 도입 계기

초기 구현에서는 Tensor가 항상 gradient를 추적하도록 만들었다. 하지만 곧 문제가 드러났다. 중간 계산에서 파생되는 모든 Tensor가 gradient를 추적하게 되면 불필요한 그래프가 생성되고, 연산 비용과 메모리 사용량이 급격히 증가한다.

특히 inference 상황에서는 gradient가 필요 없는데도 전체 computation graph가 생성되는 문제가 있었다.

이를 해결하기 위해 PyTorch와 동일한 방식으로 requires_grad 속성을 도입했다. 이 속성의 도입 이유는 다음과 같다.

1. 불필요한 gradient 추적 회피
2. 메모리 사용량 최적화
3. 사용자가 명시적으로 autodiff 활성 여부를 제어할 수 있도록 하기 위함

Tensor 생성 시 입력되는 requires_grad 값과, parent들의 requires_grad 여부를 합쳐, 결과 텐서가 gradient를 추적해야 하는지를 자동으로 결정하도록 설계했다. 이 로직은 연산마다 매우 중요한데, parent 중 단 하나라도 requires_grad=True라면 결과 역시 gradient를 추적해야 한다.

이러한 식으로 새롭게 생성된 Tensor 객체에 requires_grad와 같은 다양한 flag 작업을 하는 과정을 별도의 coroutine 을 만들어 관리하였고, 이에 대한 자세한 설계는 다음 글에서 다루어보겠다.

이 접근은 후에 구현된 네트워크 layer, optimizer 등에서 자연스럽게 작동하며 Lucid의 일관성을 유지하는 데 큰 도움이 되었다.

↩️ Backward 함수의 설계와 구현적 어려움

Tensor 클래스에서 backward() 메서드는 사실상 autodiff 엔진의 엔트리 포인트다. Lucid의 초기 backward 구현은 다음과 같은 구조를 가지고 있다.

class Tensor:
	def backward(self, keep_grad: bool = False) -> None:
    	# 1. 출력 텐서의 초기 gradient 설정
        if self.grad is None:
        	self.grad = np.ones_like(self.data)
      	
        # 2. 수동 topological sort로 역순 순회 순서 구성
        visited: set[Tensor] = set()
        topo_order: list[Tensor] = []
        stack: list[Tensor] = [self]
        
        while stack:
        	t = stack[-1]
            if t in visited:
            	stack.pop()
                topo_order.append(t)
                continue
            
            visited.add(t)
            for p in t._prev:
            	if p not in visited:
                	stack.append(p)

		# 3. topological order의 역순으로 gradient 전파
        for t in reversed(topo_order):
        	if t._op is None:
            	continue  # leaf 텐서이거나 연산 정보가 없는 경우
            
            # 각 연산 객체가 자신의 local backward를 알고 있음
            parent_grads = t._op.backward(t.grad, t, *t._prev)
            
            for parent, g in zip(t._prev, parent_grads):
            	if g is None:
                	continue
                if parent.grad is None:
                	parent.grad = g
                else:
                	parent.grad += g

			# 중간 텐서의 grad 메모리 해제
            if not (t.is_leaf or keep_grad or t.keep_grad):
            	t.grad = None

여기에서도 보이듯 backward의 핵심은 세 단계로 요약된다.

1. 현재 텐서가 scalar loss라고 가정하고 gradient를 $1$로 초기화하는 단계이다. 이는 수학적으로 $\partial L/\partial L=1$에 해당한다.

2. Computation graph를 명시적으로 위상 정렬(topological sort)하여 역순 순회 순서를 만드는 단계이다. 이 과정에서 recursive한 호출 대신 수동 stack 을 사용함으로써 깊은 그래프에서도 Python의 max recursion limit에 걸리지 않도록 했다.

3. 각 노드에서 자신의 연산 객체 _op에 정의된 local backward를 호출해 parent들로 gradient를 전파하는 단계이다. 이때 parent의 grad는 단순 대입이 아니라 += 방식으로 누적된다. 이는 수학적으로도 자연스러운 선택인데, 하나의 parent Tensor가 여러 children 연산에 동시에 관여하는 경우

   $\frac{\partial L}{\partial\mathbf{x}}=\sum_k\frac{\partial L}{\partial \mathbf{y}_k}\cdot\frac{\partial \mathbf{y}_k}{\partial \mathbf{x}}$

   와 같이 여러 경로를 따라 전달되는 gradient를 전부 더해야 올바른 미분값이 되기 때문이다.

   만약 parent.grad = grad처럼 마지막 경로의 값만 남긴다면, 실제로는 존재하는 경로 일부가 완전히 무시되어 잘못된 gradient를 얻게 된다. 그래서 Lucid에서는 처음 gradient를 받는 시점에서는 None $\rightarrow$ g로 초기화하고, 그 이후로는 항상 +=를 사용해 모든 경로의 contribution을 합산하도록 구현했다.

이 구현을 통해 나는 앞서 구상했던 "Tensor를 node로 가진 computation graph를 역순으로 따라가며 chain rule을 적용한다" 는 아이디어가 실제 코드 수준에서 어떻게 실현될 수 있는지를 검증할 수 있었다. 또한 keep_grad와 Tensor.keep_grad 플래그를 통해, 중간 텐서의 gradient를 보존할지 여부를 세밀하게 제어할 수 있게 되었다.

기본적인 철학은 "leaf Tensor의 grad만 남기고 나머지는 가능한 한 버린다" 에 가깝다.

☘️ is_leaf 플래그

is_leaf는 해당 Tensor가 사용자가 직접 생성한 입력 또는 학습 파라미터처럼, 그래프 상에서 더 이상 "연산의 결과" 로 쪼개지지 않는 말단 노드(terminal node)인지 나타낸다. 이런 텐서들의 grad는 optimizer가 실제 weight update에 사용하므로 반드시 남겨야 한다.

반대로 연산 중간에 만들어지는 non-leaf Tensor 들은 역전파가 끝난 후에는 grad가 더 이상 필요하지 않은 경우가 많다. 이때 is_leaf == False이고, global한 keep_grad=False, 그리고 개별 텐서의 t.keep_grad=False이면 해당 텐서의 grad를 과감히 None으로 되돌려 메모리를 회수한다.

다만, 디버깅 혹은 시각화 등을 염두에 두고 특정 중간 텐서의 gradient를 보고 싶을 때가 있다. 이런 경우를 위해 개별 텐서 단위로 keep_grad=True를 줄 수 있도록 했고, backward(keep_grad=True) 인자를 사용하면 한 번의 backward 동안 전체 그래프의 grad 값을 보존하도록 만들었다.

이러한 플래그 설계는 결국 메모리 효율성과 디버깅 편의성 사이의 균형점이었다. is_leaf, requires_grad, keep_grad의 조합으로 어떤 텐서의 gradient를 얼마나 오래 유지할지를 세밀하게 조절할 수 있도록 설계한 것이다.

😩 구현 과정에서의 어려움

전체적인 Tensor 구현 과정에서 마주쳤던 가장 큰 난제는, 한 번만 사용되는 단순한 연산 그래프가 아닌 복잡하게 공유된 sub-graph를 가진 네트워크 구조에서 gradient가 올바르게 누적되는지 검증하는 일이었다. 이 검증을 위해 일부러 같은 텐서를 여러 경로로 흘려보내는 실험용 computation graph들을 만들고, 수작업으로 계산한 미분 결과와 Lucid의 autodiff 엔진으로 돌아가는 backward 결과를 비교하는 과정을 반복했다.

---

🪄 Operation 코루틴의 구조로 넘어가며

Tensor의 핵심 구조인 데이터, 그래프 메타데이터, backward 로직 등이 자리 잡으면서, 연산(operation) 단위에서 생성되는 새로운 Tensor에 대한 메타데이터 초기화 과정 역시 보다 체계적인 관리가 필요해졌다.

각 연산마다

• 부모 노드 연결(_prev 설정)
• 연산 객체 등록(_op 설정)
• requires_grad 전파 규칙
• is_leaf 판정
• backward 시 parent.grad 누적 정책

등을 반복적으로 작성하는 것은 유지보수와 일관성 측면에서 매우 비효율적이다.

그래서 다음 단계에서는 이 과정을 operation-level coroutine(코루틴)으로 분리해, 모든 연산이 동일한 패턴으로 Tensor을 생성하고 그래프에 연결될 수 있도록 구조화하게 된다.

이 다음 세 번째 문서에서는 바로 이 operation 메타데이터 세팅을 자동화하는 코루틴 시스템을 중점적으로 다루고자 한다.

🧭 마무리: Tensor가 Lucid의 기반이 되기까지

Tensor의 구현은 Lucid 전체의 기반을 이루는 단계였고, 이후의 모든 모듈은 이를 통해 만들어진 autodiff 엔진 위에서 동작하게 된다. 무엇보다도 backward 구조가 완성되면서 비로소 Lucid는 "딥러닝 프레임워크" 라고 부를 수 있는 최소 조건을 만족하게 되었다.

Tensor는 단순한 배열에서 탈피하여, 데이터와 연산 정보를 가진 완전한 계산 단위로 자리 잡았고, 이 객체를 기반으로 이루어진 computation graph는 딥러닝 모델의 forward/backward를 온전히 표현하는 구조가 되었다.

이 글은 Tensor가 어떻게 등장했는지, 어떤 의사결정과 시행착오를 거쳐 이 글에서 보여준 구조로 안착했는지를 기록한 하나의 일지 이며, Lucid의 이후 발전 역시 이 기초를 기반으로 단계적으로 쌓아올려졌다.