기술 요약: 텐서 네트워크 커널 머신: 머신러닝 및 비선형 시스템 식별을 위한 JAX 프레임워크

문제 정의

높은 표현력과 계산 효율성 사이의 균형을 맞추는 비선형 모델을 개발하는 것은 머신러닝 및 비선형 시스템 식별 분야에서 여전히 중요한 과제이다. 커널 방법, 가우시안 프로세스, 심층 신경망과 같은 전통적인 접근 방식은 모델의 복잡도나 훈련 샘플 수가 증가함에 따라 모델 복잡성이나 계산 비용이 기하급수적으로 증가하는 확장성 문제를 겪는 경우가 많다. 구체적으로, 고차원 공간에서 비선형 매핑을 명시적인 특징 맵(feature map)으로 표현하는 방식은 파라미터 수의 지수적 폭발을 초래하여, 고차원 문제에 대한 직접적인 최적화를 불가능하게 만든다.

방법론

본 논문은 명시적인 비선형 특징 표현과 압축된 저계수(low-rank) 텐서 네트워크 파라미터화를 결합하여 확장성 문제를 해결하는 텐서 네트워크 커널 머신(Tensor Network Kernel Machines, TNKM) 프레임워크를 소개한다. 핵심 수학적 정식화는 비선형 함수 $f(x)$를 고차원 특징 공간에서의 선형 모델로 나타낸다:
$$f(x) = \phi(x)^\top w$$
여기서 $\phi(x)$는 비선형 특징 맵이고 $w$는 파라미터 벡터이다. $w$에 내재된 차원의 저주를 극복하기 위해, TNKM 접근 방식은 두 가지 핵심 전략을 채택한다:

1. 텐서 곱 특징 맵(Tensor-Product Feature Maps): 특징 맵 $\phi(x)$는 개별 입력 변수에 작용하는 1차원 맵들의 텐서 곱으로 구성되며, 이는 곱 커널(product kernels)로부터 자연스럽게 발생한다.
2. 저계수 텐서 분해(Low-Rank Tensor Decomposition): 전체 파라미터 벡터 $w$를 저장하는 대신, 파라미터를 CP(Canonical Polyadic) 및 TT(Tensor-Train) 형식과 같은 저계수 텐서 분해를 사용하여 표현한다. 이를 통해 파라미터 복잡도를 입력 차원에 대해 지수적에서 선형(또는 근선형)으로 감소시킨다.

저자들은 이 프레임워크를 구현하기 위해 JAX 기반의 오픈 소스 파이썬 라이브러리인 tnkm을 개발하였다. 이 라이브러리는 세 가지 주요 구성 요소를 분리하는 모듈형 아키텍처를 채택한다:

• 특징 맵(Feature Maps): 다항식(polynomial), 푸리에(Fourier), 볼테라(Volterra, 메모리 기반) 특징을 지원한다.
• 텐서 네트워크 파라미터화(Tensor-Network Parameterizations): 전용 모델 클래스(예: CPKM, TTKM)를 통해 CP 및 TT 분해를 구현한다.
• 최적화 알고리즘(Optimization Algorithms): 두 가지 상호 보완적인 전략을 제공한다:
  • 교대 최소제곱법(Alternating Least Squares, ALS): 최소제곱 문제에 대해 빠르고 결정론적인 수렴을 위해 다선형 구조를 활용한다.
  • 경사 하강 기반 최적화(Gradient-Based Optimization): 임의의 미분 가능한 손실 함수에 대한 유연성을 위해 JAX의 자동 미분과 Optax 라이브러리(예: Adam, SGD)를 활용한다.

주요 기여

본 논문은 네 가지 주요 기여를 명시한다:

1. 통합 JAX 프레임워크: 서로 다른 특징 맵, 텐서 네트워크 백본(CP 및 TT), 최적화 전략을 단일 인터페이스 아래 통합한 최소 의존성 오픈 소스 라이브러리(tnkm)를 제공한다.
2. 확장 가능한 훈련 구현: 구조 인지형 ALS 최적화를 일반적인 경사 하강 기반 방법과 함께 통합하여, 회귀 및 더 넓은 학습 작업 모두에 적용할 수 있도록 한다.
3. 모듈형 소프트웨어 아키텍처: 구성 요소(특징 맵, 모델 클래스, 옵티마이저)를 체계적으로 교체할 수 있는 설계를 제공하며, 비선형 시스템 식별에 특화된 데이터 전처리, 시뮬레이션 및 평가 유틸리티를 제공한다.
4. 포괄적인 실험적 평가: 비선형 회귀 및 시스템 식별 작업에 대한 엄격한 벤치마킹을 통해 효율적인 훈련 시간과 함께 경쟁력 있는 정확도를 입증하였다.

실험 결과

저자들은 두 가지 측면에서 tnkm을 평가하였다:

1. 최적화 방법 비교 (Airfoil Self-Noise 데이터셋):

• ALS vs. Adam: ALS는 1.16초 만에 검증 MSE 0.102에 빠르게 수렴하여 커널 리지 회귀(KRR) 베이스라인(0.15)에 근접하였다. 반면, Adam은 더 높은 오차(0.145)에 도달하기 위해 훨씬 더 많은 시간(미니 배치 기준 3.25초)과 에포크(epoch)가 필요했다.
• 결론: ALS는 최소제곱 TNKM 문제에 대해 우수한 효율성을 제공하는 반면, 경사 하강 기반 방법은 비 최소제곱 목적 함수를 위한 필수적인 유연성을 제공한다.

2. 비선형 시스템 식별 벤치마크:
프레임워크는 Coupled Electric Drives, Cascaded Tanks, Silverbox의 세 가지 표준 벤치마크에 대해 테스트되었다.

• Coupled Electric Drives: TNKM은 가장 성능이 좋은 방법들(GPNARX, dynoNet)과 대등한 RMSE를 달면서도, 경쟁 방법들이 수 초에서 수 분이 걸리는 데 반해 1초 미만의 훈련 시간을 기록하였다.
• Cascaded Tanks: TNKM은 블랙박스 접근 방식 중 경쟁력 있는 정확도(RMSE 0.347)를 제공하였으며, SUBNET 및 PNARX와 같은 방법들에 비해 훈련 속도 면에서 압도적이었다(1~60분 대비 < 0.1초).
• Silverbox: TNKM은 선도적인 블랙박스 방법들(GPNARX, SUBNET)과 대등한 성능을 보이면서도 1~60초의 훈련 시간을 기록하였다. 물리 정보 기반 접근 방식(NL-LFR)의 정확도에는 미치지 못했으나, 효과적인 장기 예측 안정성을 입증하였다.

의의 및 주장

본 논문은 tnkm이 텐서 네트워크 기반 학습 방법의 재현 가능한 개발 및 적용을 용이하게 한다고 주장한다. 특징 맵, 텐서 파라미터화, 최적화 알고리즘을 분리함으로써, 이 라이브러리는 기본 워크플로우를 변경하지 않고도 TNKM 구성을 체계적으로 탐색할 수 있게 한다.

저자들은 TNKM을 고전적인 커널 방법과 심층 신경망 사이의 중간 단계 접근법으로 위치시킨다. 커널 방법과 마찬가지로 TNKM은 명시적인 특징 표현을 사용하여 더 큰 제어력과 사전 지식의 통합을 제공하지만, 확장성을 보장하기 위해 전체 파라미터 벡터를 압축된 텐서 네트워크로 대체한다. 신경망과 달리, TNKM은 데이터로부터 특징을 직접 학습하지 않으며, 저자들은 이를 수동적인 특징 선택이 필요한 한계점으로 인정한다.

본 연구의 의의는 이론적인 텐서 네트워크 방법과 실제 세계의 비선형 시스템 식별 사이의 간극을 메우는 실용적이고 확장 가능한 소프트웨어 기반을 제공한다는 점에 있으며, 이러한 모델이 낮은 계산 비용을 유지하면서도 높은 예측 정확도를 달성할 수 있음을 입증하였다.