CodeJeNN: A simple C++ neural network generator for physics applications

기술 요약: CodeJeNN

문제 정의
머신러닝은 난류 모델링 및 대리 구성 모델(surrogate constitutive models)과 같은 작업에서 계산 유체 역학(CFD)을 포함한 수치 해석 방법론을 가속화하는 데 상당한 잠재력을 입증해 왔다. 그러나 Python 기반 머신러닝 라이브러리(예: TensorFlow, PyTorch, Keras)를 고성능 컴파일형 C++ 또는 Fortran 솔버에 통합하는 과정에서 결정적인 병목 현상이 존재한다. 이러한 불일치는 계산 오버헤드를 유발하고, 기존 물리 프레임워크 내에 신경망을 배포하는 것을 복잡하게 만든다. 또한, 프레임워크가 제공하는 C++ API에 의존하는 것은 텐서 연산을 외부 라이브러리로 강제하여 기존의 병렬화 및 벡터화 전략과 충돌할 수 있으므로 최적화 능력을 제한한다. RoseNNa와 같이 Fortran 및 ONNX 모델을 위한 기존 솔루션들이 존재하지만, 외부 런타임 의존성 없이 원활한 통합을 촉진하기 위해 Keras 모델에 특화된 자립형(self-contained), 의존성 없는 C++ 코드를 생성하는 도구에 대한 필요성이 여전히 남아 있다.

방법론
저자들은 Keras 모델과 컴파일된 물리 솔버 사이의 간극을 메우는 C++ 신경망 생성기인 CodeJeNN을 제시한다. 워크플로우는 다음과 같이 작동한다:

1. 모델 학습: 사용자는 Keras를 사용하여 신경망(MLP 또는 CNN)을 학습시키고, 아키텍처와 파라미터를 .keras 또는 .h5 파일로 저장한다.
2. 코드 생성: CodeJeNN은 모델을 로드하여 파라미터를 Python 딕셔너리로 추출하고 모델 로직을 재구성한다.
3. C++ 합성: 도구는 자립형 C++ 헤더 파일을 생성한다. 주요 설계 원칙은 다음과 같다:
   • 인라이닝(Inlining): 함수 호출 오버헤드를 최소화하기 위해 활성화 함수 및 전파 함수를 공격적으로 인라이닝한다.
   • 의존성 감소: 생성된 코드는 표준 C++ 라이브러리(예: <iostream>, <array>, <cmath>)만을 사용하여 LAPACK이나 Eigen과 같은 외부 선형 대수 라이브러리의 필요성을 제거한다.
   • 컴파일 타임 최적화: constexpr, static, inline 키워드를 사용하여 배열 크기와 연산이 컴파일 타임에 결정되도록 함으로써 캐시 활용도를 최적화한다.
   • 확장성: 소스 구조는 지원되지 않는 아키텍처에 대해 최소한의 토큰 소비만으로 LLM이 로직을 생성할 수 있도록 설계되어, 비침습적인 프레임워크 확장을 용이하게 한다.

이 도구는 광범ening한 범위의 Keras 레이어(Dense, Convolutional, Pooling, Normalization 등)와 활성화 함수(ReLU, GELU, Swish 등)를 지원하며, Python 스크립트에 의해 직렬화되고 정의된 사용자 정의 커스텀 활성화 함수도 포함한다.

주요 기여

• 자립형 코드 생성: CodeJeNN은 필요한 모든 가중치, 편향(bias) 및 로직을 포함하는 단일 C++ 헤더 파일을 생성하여, 외부 런타임 의존성 없이 기존 솔버에 직접 포함될 수 있도록 한다.
• 성능 최적화: 외부 API 호출을 피하고 공격적인 인라이닝을 활용함으로써, 생성된 코드가 로컬 하드웨어에서 효율적으로 실행되도록 설계되었다.
• 디버깅 기능: 이 도구는 추출된 모델을 Keras의 model.summary()와 비교하는 디버깅 API를 포함하며, 활성화 시 추론 과정 중의 중간 레이어 출력을 출력하여 로직 또는 대수적 오류를 찾아내는 데 도움을 준다.
• LLM을 통한 확장성: 코드베이스는 LLM을 사용하여 새로운 아키텍처 지원을 추가할 때 발생하는 토큰 비용과 노력을 최소화하도록 구조화되었다.

결과
본 논문은 두 가지 주요 벤치마크를 통해 CodeJeNN을 검증한다:

1. 추론 속도 벤치마크: Apple M1 Pro에서 TensorFlow 및 PyTorch 백엔드(eager 및 JIT 컴파일 모드 모두)에서 실행되는 Keras와 CodeJeNN을 비교하는 테스트를 수행하였다.
   • 다층 퍼셉트론(MLP)의 경우, CodeJeNN은 표준 TensorFlow 대비 25배, 표준 PyTorch 대비 38배 빠른 속도를 달了했다.
   • JIT 컴파일된 Keras 모델과 비교했을 때도, CodeJeNN은 약 2.5배(TensorFlow) 및 9배(PyTorch) 더 빨랐다.
   • 합성곱 신경망(CNN)의 경우, 상당한 속도 향상(표준 Keras 대비 18x~26x)을 보였으나, 모델이 커짐에 따라 캐시 제한으로 인해 연산 집약적(compute-bound) 단계에서 메모리 집약적(memory-bound) 단계로 전환되면서 성능이 저하되었다.
2. CFD 응용 (점도 모델링): 수소-공기 혼합층의 점도 계산을 위해 Wilke의 혼합 규칙을 대체하도록 신경망을 OpenFOAM 기반 솔버(DetonationFOAM)에 임베딩하였다.
   • 10,000개의 상태(states)로 학습된 네트워크가 솔버의 C++ 헤더에 직접 삽입되었다.
   • 시뮬레이션 결과, 신경망 대리 모델은 네이티브 Wilke 모델과 비교하여 속도장(velocity field)에서 최대 **0.25%**의 오차만을 도입하였으며, 전체적인 흐름 진화 및 불안정성 발달을 보존하였다.

의의 및 주장
저자들은 CodeJeNN이 물리 연구자들이 고성능 솔버에 신경망을 통합하는 장벽을 크게 낮춘다고 주장한다. Raw 인라인 소스 코드를 생성함으로써, 복잡한 인프라 구축이나 외부 텐서 라이브러리에 대한 의존성을 제거하며, 이는 가용 가능한 모든 하드웨어 자원을 활용하는 코드에 매우 중요하다. 또한, LLM을 사용하여 인프라 코드를 반복적으로 생성하는 비효율성을 해결하기 위해 전용의 반복 가능한 파이프라인을 제공한다. 본 논문은 실시간 및 하드웨어 제약 조건을 충족하기 위해 모델이 컴팩트하고 효율적이며 외부 의존성이 없어야 하는 과학 계산 분야에서 CodeJeNN을 실용적인 솔루션으로 포지셔닝한다. 향러 연구에는 Spektral을 통한 GNN(Graph Neural Network) 지원 통합과 추가적인 가속을 위한 저수준 라이브러리(예: Eigen) 내보내기 탐색이 포함된다.