이번 주 Colab에서 JAX 실행을 확인하자마자, Equinox 버전과 사용 가능한 디바이스를 출력하는 장면이 나온다.

또한 Linear 모듈을 만들고 PyTree leaves와 구조를 바로 확인하는 흐름이 이어진다.

여기서 모델이 PyTree로 취급된다는 전제가 코드 점검 단계에서 먼저 드러난다.

Equinox 환경 세팅과 PyTree 기반 모듈 구성

튜토리얼은 Equinox를 JAX 위에 올린 신경망 라이브러리로 두고, JAX, Equinox, Optax, Jaxtyping, Matplotlib 등 필수 라이브러리를 설치·임포트하는 단계로 시작한다.

그 다음 JAX와 Equinox 버전을 출력하고, 사용 가능한 디바이스 목록을 확인해 Colab 실행 준비 상태를 검증한다.

이후 eqx.Module이 모델을 PyTree로 취급해 파라미터 처리, 변환, 직렬화가 단순하고 명시적으로 느껴진다고 설명한다.

filter_jit·filter_grad로 배열 파트만 학습에 참여시킴

튜토리얼은 Linear 레이어를 포함한 간단한 모듈을 만든 뒤, Conv1dBlock을 통해 static field와 learnable layer가 함께 동작하는 예를 든다.

이후 드롭아웃을 포함한 MLP를 구성하고, filter_jit으로 forward pass를 컴파일하되 모델에 비배열(non-array) 필드가 있어도 허용하도록 한다.

반면 filter_grad는 실제로 학습에 참여해야 하는 배열 leaf만 골라 그래디언트를 계산한다.

합성 데이터로 forward를 돌리고 그 결과로 그래디언트를 평가하면서, Equinox가 모델 정의와 미분 가능한 계산을 JAX 친화적으로 연결한다고 보여준다.

PyTree 분할·tree_at 불변 업데이트·BatchNorm 상태 처리

PyTree 조작 유틸리티 파트에서는 모델을 배열과 비배열로 분할하고, 첫 레이어를 고정하기 위한 trainable filter를 만든다.

그 다음 tree_at으로 특정 파라미터에 대해 불변(immutable) 업데이트를 수행하며, 전체 모델을 다시 쓰지 않는 방식을 사용한다.

이후 stateful layer로 BatchNorm 기반 모델을 정의하고, 모델과 state를 각각 생성한 뒤 배치 학습 스타일 패스를 돌려 업데이트된 state 정보를 반환받는다.

여기서 주목할 점은, 학습 루프가 단순히 가중치만 바꾸는 형태가 아니라 state까지 함께 다룬다는 점이다.

잔차 블록 학습, warmup cosine 스케줄, 직렬화·make_jaxpr 디버깅

튜토리얼은 noisy sine 회귀 작업을 위한 residual block과 ResNetMLP 모델을 정의해 더 깊은 end-to-end 예제를 만든다.

합성 학습/검증 데이터셋을 생성하고, 모델을 초기화한 뒤 warmup cosine learning rate schedule을 구성한다.

옵티마이저 상태는 모델의 배열 leaf만 사용해 준비하고, jitted train_step과 evaluate 함수를 정의해 학습과 검증의 핵심 메커니즘을 제공한다.

학습 루프는 여러 epoch 동안 데이터를 셔플하고 미니배치를 처리하며, 학습 손실과 검증 손실을 시간에 따라 추적한다.

그 다음 Equinox 유틸리티로 학습된 모델을 직렬화하고, 동일한 스켈레톤 모델을 재구성한 뒤 deserialization이 가중치를 올바르게 복원하는지 확인한다.

또한 jax.make_jaxpr로 컴파일된 계산 그래프를 검사해, JAX에서 실행되는 trained Equinox 모델의 동작을 디버깅·인트로스펙션 관점에서 본다.

개발자가 얻는 핵심은, PyTree 기반 모듈 정의부터 filter_jit·filter_grad의 선택적 컴파일/미분, BatchNorm state 처리, 직렬화 검증, make_jaxpr 그래프 점검까지 한 흐름으로 연결된다는 점이다.