Skip to content

Julia + ML?

작년에 neural ODE4에 관해 공부하면서 Julia에 대해 관심이 생겼다. Julia를 사용하기 위해 이런 저런 노력을 했지만 결국에 포기하고 말았는데 이유를 글로 정리해 남겨두면 좋을 것 같다.

Why Julia?

연구 과정 중 Python으로 작성한 코드를 정리하다보면 (주로 비효율적으로 작성된 부분을 개선하는 작업) 조금 과장을 보태어 "이럴 거면 C로 다시 작성하는게 낫지 않나?"라는 물음에 도달한다. 이를 two language problem이라고 부르기도 하는데, two language problem이란 쉬운 Matlab, Python을 사용해서 빠르게 알고리즘을 테스트한 후, 잘 작동한다 싶으면 C나 Fortran으로 performance optimization을 하는 것을 말한다1. 하지만 C나 Fortran은 내게 너무 어려웠고, 생산성 유지를 위해 상대적으로 느린 Python에 머물러 있을 수밖에 없었다. 근데 연구를 진행하다 보면 남이 만들어둔 코드를 베이스로 해서 무언가를 더 얹는 경우가 많은데, 이 과정이 반복되면 이상적인 코드와 비교하여 굉장히 느릴 수 있다는 생각이 들었고, Julia가 작성하기 쉬우면서 동시에 높은 잠재력을 가졌다는 글5을 보고 관심이 생겼다.

Julia는 two language problem을 해결할 수 있을 것처럼 보인다. 벤치마크를 보면 최대 성능은 Fortran이나 Rust만큼 빠르다. 또한 함수를 작성할 때 type을 지정해 줄 필요가 없기 때문에 생산성도 충분히 나온다.

Note

아무렇게나 작성한 함수가 이상적인 C코드만큼 빠르게 동작하는 것은 아니다. Type stability와 heap allocation을 최소화해야 jit 컴파일 이후 빠르게 동작한다.

함수 이름이나 변수 이름으로 특수문자를 사용할 수 있기 때문에 readability 측면에서도 훌륭한 잠재력을 가졌다. 예를 들어 다음과 같은 표현이 가능하다.

function ∂t(f, t, ϵ=1e-3)
  return (f(t + ϵ) - f(t - ϵ)) / (2 * ϵ)
end

f(t) = t^2

∂t(f, 0.5) # 0.999999999999987

Why not Julia + ML?

이러한 훌륭한 장점들에도 불구하고, 사실 내가 하는 연구 (주로 physics-informed machine learning)에서 Julia를 사용하는 것은 쉽지 않았다. 가장 큰 문제점은 에러 메시지, 그리고 관련 ecosystem의 documentation의 불친절함이다. 먼저 에러 메시지의 경우, Python은 한줄 한줄 실행하는 특성 덕분에 몇 번째 줄에서 에러가 발생했는지 짚어주기 때문에 매우 편리하다. 반면 Julia 함수는 호출되면 LLVM 컴파일러를 거쳐 효율적인 표현으로 변환되기 때문에 어디서 에러가 발생했는지 알아내기가 상대적으로 어렵다.

불친절한 documentation은 여러 이유가 있을 것 같은데, 아무래도 큰 회사에서 관리하지 않아서 그런 것 같다. Google의 JAX, Meta의 Torch는 훌륭한 사용자들을 자사 상품의 연구/개발 인력으로 고용할 유인이 있기 때문에 사용자 풀을 늘리거나 유지하기 위해 문서를 친절하게 관리할 것이다. 하지만 Julia의 machine learning ecosystem은 규모가 작은 회사 (Flux.jl - Julia computing) 혹은 학계 (Lux.jl - MIT CSAIL)에서 개발되었다. 인력 측면에서 빅테크 기업들에 비해 큰 차이가 날 수밖에 없다. 또한 Julia는 상대적으로 젊은 언어이기 때문에 사용자들이 프로그래밍에 익숙한 경우가 많고, 따라서 documentation이 불친절하면 그냥 소스 코드를 뜯어보는데 주저함이 없는 경우가 많았을 것이다.

다른 이유로는,

  • 예전에는 jit 컴파일의 initial overhead가 너무 심해서 간단한 그림을 그리는데 1분 넘게 기다려야 하는 경우가 있었다. v1.10이 되면서 해결된 것으로 알고 있다.
  • 내가 주로 했던 physics-informed neural network (PINN) 연구에서는 optimization step 한 번마다 automatic differentiation (AD)을 여러번 사용하는 경우가 많은데, 이런 nested AD는 비교적 최근에 지원하기 시작했다. PINN 이외에도 Neural ODE를 학습할 때 Jacobian의 norm을 regularization 하는 테크닉들이 있는데, Julia로 이를 작성하려면 PyTorch에 비해 고생을 좀 해야 할 것이라고 생각한다.
  • 큰 인공신경망이 연구에 포함되는 순간 대부분의 overhead는 GPU 속도에 의존하기 때문에 Julia를 사용할 유인이 없다. (사실 이는 JAX도 마찬가지인 것 같은데..?)
  • Performance optimization은 또 다른 이야기였다. QR decomposition 같은 경우에 Gram-Schmidt process나 Householder reflection이나 목표는 같지만 implementation의 형태는 꽤 다르고, memory cost를 줄이기 위해 allocation을 줄이는 최적화된 코드는 더 다르다. 즉, 같은 프로그래밍 언어에 머물러 있더라도 최적화된 코드는 난해하고 어려울 가능성이 높다. 그렇다면 어떠한 의미에서는 two language problem은 아직 풀리지 않았다고 할 수 있지 않을까?

지금까지는 ASML, national lab에서 Julia를 사용하는 사람들을 볼 수 있었다. CPU 계산이 많은 분야에서는 훌륭한 포텐셜이 있지만, 아직은 머신 러닝 측면에서는 아닌 것 같았다. 사용자들이 더 많아지고 충분한 인력을 가진 그룹이 뛰어든다면 그제서야 편리하게 사용할 수 있지 않을까 생각한다.

P.S.

  • Julia의 단점 (에러 메시지)과 Julia ML 생태계의 단점 (premature, 문서, nested AD), 머신 러닝 연구의 특성 (GPU 속도에 의존), 그리고 performance optimization의 어려움을 구분해서 생각해야 한다.
  • 사실 인터넷에 이 토픽에 관해서 엄청나게 많은 글들이 있다23. 거기에 하나를 더하는게 좋은 일인지는 모르겠지만, 적어도 내 생각을 정리하는데는 좋을 것 같아 작성했다.

  1. 사실 Jake VanderPlas의 블로그 포스트를 보면 Numba를 사용해서 Numpy로 작성된 코드의 퍼포먼스를 C와 비슷한 수준으로 끌어올릴 수도 있다. ↩

  2. https://kidger.site/thoughts/jax-vs-julia/ ↩

  3. https://forem.julialang.org/jacobusmmsmit/thoughts-on-using-jax-and-julia-2f0b ↩

  4. Chen, R.T., Rubanova, Y., Bettencourt, J., & Duvenaud, D.K. (2018) Neural ordinary differential equations. Advances in neural information processing systems, 31. ↩

  5. Bezanson, J., Edelman, A., Karpinski, S., & Shah, V.B. (2017) Julia: A fresh approach to numerical computing. SIAM review, 59, 65–98. ↩