JAX로 WebGL 위에 레이마칭 렌더러를 만든다고? 머신러닝 프레임워크의 의외의 활용법

머신러닝 프레임워크로 3D 그래픽을 렌더링한다?

JAX라고 하면 보통 딥러닝이나 수치 계산을 떠올리잖아요. Google이 만든 고성능 수치 연산 라이브러리로, NumPy와 비슷한 인터페이스를 제공하면서도 GPU/TPU 가속과 자동 미분을 지원하는 도구인데요. 그런데 이 JAX를 가지고 레이마칭(Ray Marching) 기반 3D 렌더러를 만들고, 그것도 WebGL 위에서 브라우저로 돌리는 프로젝트가 공개됐어요. Benoit Paris라는 개발자가 진행한 이 실험적 프로젝트인데, JAX의 가능성을 정말 색다른 각도에서 보여주고 있어서 한번 자세히 살펴볼게요.

레이마칭이 뭔데?

먼저 레이마칭이 뭔지부터 이야기해볼게요. 3D 그래픽을 화면에 그리는 방법에는 여러 가지가 있는데, 가장 흔한 건 삼각형 메시(mesh)를 이용해서 물체를 표현하고 래스터라이제이션(rasterization)으로 화면에 찍는 거예요. 게임 엔진에서 주로 쓰는 방식이죠.

레이마칭은 이것과 완전히 다른 접근이에요. 카메라에서 화면의 각 픽셀 방향으로 가상의 광선(ray)을 쏘고, 그 광선을 조금씩 전진시키면서 "물체에 닿았나?"를 확인하는 거예요. 물체를 삼각형으로 표현하는 게 아니라 수학 함수(SDF, Signed Distance Function)로 표현해요. 이게 뭐냐면, 공간의 어떤 점이 주어졌을 때 "가장 가까운 물체 표면까지의 거리"를 반환하는 함수예요. 이 거리가 0에 가까워지면 "물체에 닿았다"고 판단하는 거죠.

이 방식의 장점은 아주 복잡한 형상을 수학적으로 간결하게 표현할 수 있다는 거예요. 예를 들어 구, 상자, 토러스 같은 기본 도형을 합치거나 빼거나 교차시키는 연산을 SDF 함수 조합만으로 할 수 있거든요. Shadertoy 같은 플랫폼에서 레이마칭으로 만든 놀라운 시각 효과들을 본 적 있을 텐데, 바로 이 원리예요.

JAX가 여기서 왜 등장하는 걸까

레이마칭의 핵심은 결국 대량의 수학 연산을 병렬로 처리하는 것이에요. 화면의 모든 픽셀에 대해 독립적으로 광선을 쏘고 거리 함수를 반복 계산해야 하니까요. 그런데 이거, 어디서 많이 본 패턴 아닌가요? 맞아요, 딥러닝에서 텐서 연산을 GPU로 병렬 처리하는 것과 본질적으로 같은 구조예요.

JAX는 바로 이 지점에서 빛을 발해요. JAX의 vmap이라는 기능을 쓰면 하나의 픽셀에 대한 레이마칭 로직을 작성해놓고, 그걸 자동으로 전체 픽셀에 대해 병렬화할 수 있거든요. 수동으로 GPU 셰이더를 작성할 필요 없이, 파이썬으로 직관적인 코드를 쓰면 JAX가 알아서 벡터화해주는 거예요.

여기에 더해 JAX의 자동 미분(autodiff) 기능도 활용할 수 있어요. SDF 함수를 미분하면 물체 표면의 법선 벡터(normal vector)를 구할 수 있는데, 이건 조명 계산에 필수적인 값이에요. 보통은 이 법선을 수치적으로 근사하거나 수동으로 수식을 유도해야 하는데, JAX를 쓰면 jax.grad로 자동 계산이 가능한 거죠. 이게 정말 강력한 부분이에요.

WebGL 위에서 돌린다는 건 어떤 의미일까

이 프로젝트의 또 다른 흥미로운 점은 결과물이 웹 브라우저에서 실행된다는 거예요. JAX 코드를 WebGL 백엔드로 컴파일해서 브라우저의 GPU를 활용하는 방식인데요. 이게 가능한 이유는 JAX가 내부적으로 XLA(Accelerated Linear Algebra)라는 컴파일러를 사용하고, 이 컴파일러가 다양한 하드웨어 타겟으로 코드를 생성할 수 있기 때문이에요.

파이썬으로 작성한 레이마칭 코드가 최종적으로 WebGL 셰이더처럼 동작한다고 생각하면 돼요. 별도의 앱 설치 없이 브라우저만 있으면 3D 렌더링 결과를 볼 수 있으니까, 데모나 프로토타이핑에 아주 편리한 거죠.

기존 방식들과 비교하면

전통적으로 레이마칭을 구현하려면 GLSL이나 HLSL 같은 셰이더 언어로 직접 코드를 작성해야 했어요. Shadertoy가 대표적인 플랫폼인데, 문법이 C 스타일이고 디버깅이 어렵기로 유명하죠. 반면 JAX를 쓰면 파이썬의 편안한 환경에서 코드를 작성하고, NumPy와 유사한 API를 쓸 수 있으니 진입 장벽이 훨씬 낮아요.

비슷한 맥락에서 최근 미분 가능한 렌더링(Differentiable Rendering) 분야가 주목받고 있는데요. NVIDIA의 Kaolin, Google의 TensorFlow Graphics, 그리고 PyTorch3D 같은 프로젝트들이 있어요. 이들은 3D 렌더링 과정 전체를 미분 가능하게 만들어서, 렌더링 결과로부터 역으로 3D 장면의 파라미터를 학습할 수 있게 하는 기술이에요. NeRF(Neural Radiance Fields)가 바로 이 기술의 유명한 응용 사례죠.

JAX 기반 레이마칭은 이런 미분 가능 렌더링 연구의 프로토타이핑 도구로도 매력적이에요. JAX의 자동 미분이 렌더링 파이프라인 전체에 자연스럽게 적용되니까, 역방향으로 그래디언트를 전파하는 실험을 쉽게 해볼 수 있는 거죠.

한국 개발자에게 주는 시사점

이 프로젝트가 던지는 메시지는 "도구의 경계를 고정관념에 가두지 말자"는 거예요. JAX는 머신러닝 프레임워크지만, 본질적으로는 고성능 수치 연산 도구예요. 행렬 곱셈, 벡터 연산, GPU 가속, 자동 미분 같은 기능은 머신러닝뿐 아니라 물리 시뮬레이션, 컴퓨터 그래픽스, 로보틱스, 금융 공학 등 다양한 분야에서 쓸 수 있거든요.

실무적으로 당장 쓸 수 있는 상황은 아니더라도, 몇 가지 배울 점이 있어요. 첫째, JAX의 vmap, jit, grad 같은 함수 변환(function transformation) 패턴은 병렬 처리 코드를 작성할 때 굉장히 유용한 사고 방식을 제공해요. 둘째, 이 프로젝트를 따라해보면서 셰이더 프로그래밍과 SDF 기반 모델링을 파이썬으로 배울 수 있어요. 게임이나 시각화 관련 업무를 하는 분들에게 좋은 학습 자료가 될 수 있죠. 셋째, 미분 가능 렌더링이나 NeRF 같은 분야에 관심 있는 분이라면, JAX 기반으로 이런 실험을 시작하는 것도 좋은 선택이에요.

마무리

머신러닝 프레임워크의 강력한 수치 연산 능력은 ML 모델 훈련에만 국한되지 않아요. JAX로 레이마칭 렌더러를 만드는 이 프로젝트는, 도구를 바라보는 시야를 넓히면 정말 창의적인 결과물이 나올 수 있다는 걸 보여주는 사례예요.

여러분은 원래 용도와 다르게 활용해본 개발 도구가 있나요? 혹은 JAX나 PyTorch 같은 ML 프레임워크를 ML이 아닌 다른 목적으로 써본 경험이 있다면 이야기를 나눠봐요!