Performance Benchmarking of Tensor Trains for accelerated Quantum-Inspired Homogenization on TPU, GPU and CPU architectures

이 논문은 JAX를 사용하여 양자 영감 기반의 SFFT 기반 균질화 알고리즘을 적응 및 가속화함으로써 CPU, GPU 및 TPU 상의 텐서 트레인 연산을 벤치마킹하며, 이를 통해 기존 GPU 기반 FFT 방식으로는 불가능했던 3억 개에서 700억 개의 격자점을 아우르는 고해상도 다중 스케일 시뮬레이션을 성공적으로 가능하게 한다.

핵심

거대한 문제: 너무 많은 데이터, 너무 적은 공간

당신이 복잡한 재료(예: 첨단 금속 합금이나 복합 재료)가 스트레스를 받을 때 어떻게 반응하는지 이해하려고 노력 중이라고 상상해 보세요. 이를 위해 과학자들은 재료의 미세한 내부 구조를 관찰하기 위해 "현미경"을 사용합니다.

과거에 이 현미경들은 작고 다루기 쉬운 사진들을 제공했습니다. 하지만 새로운 기술은 수백억 개의 작은 픽셀(복셀이라고 불리는)을 포함하는 초고해외상도 이미지를 우리에게 선사하고 있습니다.

문제는 이러한 거대한 이미지를 전통적인 방식으로 계산하려고 하면, 마치 종이 봉투에 산더미 같은 모래를 담으려는 것과 같다는 점입니다. 컴퓨터는 메모리가 부족해지거나(봉투가 찢어짐), 계산하는 데 시간이 너무 오래 걸려 결과가 도착했을 때는 이미 쓸모가 없게 됩니다.

해결책: "양자 영감(Quantum-Inspired)" 압축

저자들은 **텐서 트레인(Tensor Trains, TT)**이라는 수학적 기법을 사용하여 이 데이터를 처리하는 새로운 방법을 제안합니다.

재료의 데이터를 수십억 개의 작은 블록으로 만들어진 거대한 3D 루빅스 큐브라고 생각해 보세요.

• 기존 방식 (FFT): 모든 블록을 하나하나 개별적으로 살펴보며 문제를 해결하려고 합니다. 이는 데이터를 저장하기 위한 거대한 창고와 숫자를 계산하기 위한 슈퍼컴퓨터를 필요로 합니다.
• 새로운 방식 (Tensor Trains): 모든 블록을 일일이 저장하는 대신, 큐브에 일정한 패턴이 있다는 것을 깨닫습니다. 당신은 블록들이 어떻게 연결되는지 알려주는 몇 개의 "설명서"(코어라고 불림)만을 저장함으로써 전체를 설명할 수 있습니다. 이것은 4K 영화를 화질 저하 없이 아주 작은 파일로 압축하는 것과 같습니다.

이 방법은 실제 양자 컴퓨터가 아닌 일반 슈퍼컴퓨터에서 수학 문제를 풀더라도, 양자 물리학의 기술(양자 푸리에 변환)에서 아이디어를 빌려왔기 때문에 "양자 영감" 방식이라고 불립니다.

실험: 누가 가장 빠른 러너인가?

저자들은 이 새로운 "압축된" 방식이 현대의 컴퓨터 칩에서 얼마나 빠르게 작동하는지 확인하고 싶었습니다. 그들은 세 가지 다른 유형의 하드웨어를 테스트했습니다:

1. CPU: 컴퓨터의 표준 뇌 (믿음직하고 다용도로 쓰이는 일꾼과 같음).
2. GPU: 그래픽 및 병렬 처리를 위해 설계된 칩 (함께 일하는 10,000마리의 개미 군단과 같음).
3. TPU: 구글이 AI를 위해 특별히 만든 전문 칩 (특정 트랙만을 위해 제작된 포뮬러 원 레이싱 카와 같음).

그들은 이 "압축된" 수학을 이 칩들에서 실행하기 위해 새로운 엔진(JAX라는 소프트웨어 도구 사용)을 만들었고, 속도를 측정했습니다.

결과: 경주에 따라 답은 달라진다

논문은 단 하나의 "승자"는 없다는 것을 밝혀냈습니다. 그것은 문제의 크기와 수행하는 수학의 종류에 달려 있습니다.

• 거대하고 병렬적인 작업의 경우 (GPU의 승리): 수학 계산이 한 번에 수백만 개의 단순한 계산을 수행하는 것(예: 거대한 리스트를 더하기)을 포함할 때, GPU가 가장 빨랐습니다. GPU는 엄청난 데이터셋을 처리하며 매우 아름답게 확장됩니다.
• 작거나 더 복잡한 작업의 경우 (TPU의 승리): 작업을 나누기 어려운 특정 유형의 수학의 경우, TPU가 놀라울 정도로 효율적이었으며, 종종 CPU를 이기고 때로는 GPU를 앞지르기도 했습니다.
• CPU: 가장 느렸지만 가장 안정적이었습니다. 데이터가 너무 커져도 멈추지 않았으며, 가속기(GPU, TPU)들이 메모리 부족으로 멈추는 동안에도 안정성을 유지했습니다.

매트릭스의 오류:
저자들은 TPU에서 특정 문제를 발견했습니다. 매우 높은 정밀도의 큰 숫자로 특정 유형의 복잡한 수학(SVD라고 불림)을 수행할 때, TPU가 혼란을 느껴 제대로 작동하지 않는 현상이 발생했습니다. 이를 해결하기 위해, 그들은 TPU만을 위해 약간 더 느리지만 더 안정적인 "백업 플랜"(Polar Decomposition)을 사용해야 했습니다.

최종 결론: 한계를 돌파하다

이 논문에서 가장 흥期待되는 부분은 이 새로운 설정으로 무엇을 달성했는가 하는 점입니다.

그들은 700억 개의 그리드 포인트를 가진 데이터셋에 대해 균질화 시뮬레이션을 성공적으로 실행했습니다.

• 주의 사항: 기존의 최선책들(표준 FFT 사용)은 이 작업을 아예 수행할 수 없습니다. 이 크기에 도달하기도 전에 메모리가 바닥나 버립니다.
• 돌파구: 이 가속기들 위에서 "압축된" 텐서 트레인 방식을 사용함으로써, 그들은 이전에는 불가능했던 문제들을 해결할 수 있었습니다.

요약

이 논문을 세 가지 자동차(CPU, GPU, TPU)에 탑재된 새로운 연료 효율적인 엔진(Tensor Trains)에 대한 시승 테스트라고 생각하세요.

• 그들은 이 새로운 엔진이 기존 엔진보다 훨씬 더 멀리 갈 수 있음(더 큰 데이터를 처리할 수 있음)을 증명했습니다.
• GPU는 긴 직선 고속도로 주행(거대한 병렬 데이터)에 가장 좋은 자동차라는 것을 발견했습니다.
• TPU는 특정 기술적 트랙에 훌륭하지만, 고정밀 수학에서 몇 가지 특이한 점이 있다는 것을 발견했습니다.
• 가장 중요한 것은, 이 새로운 엔진을 통해 우리는 예전에 완전히 막혀 있던 "교통 체증"(거대한 데이터셋)을 뚫고 지나갈 수 있게 되었다는 것입니다.

기술 요약

기술 요약: TPU, GPU 및 CPU 아키텍처 상에서의 양자 영감형 균질화를 위한 텐서 트레인 성능 벤치마킹

문제 정의
최근 고해상도 CT 이미징의 발전은 초고해상도 미세구조 데이터셋(수백억 개의 복셀에 달함)을 생성하였으며, 이는 전통적인 균질화 방식에 도전 과제를 제기하고 있다. 최첨단 고속 푸리에 변환(FFT) 기반 균질화 기법은 중간 규모의 데이터셋에는 효과적이지만, 메모리 점유율과 계산 비용이 $O(dN^d \log N)$으로 스케일링되어 산업 규모의 문제에는 비효율적이다. GPU 및 TPU와 같은 하드웨어 가속기가 강력한 연산 능력을 제공하지만, 고해상도 데이터의 극심한 메모리 요구량은 종종 이들의 용량을 초과한다. 양자 푸리에 변환(QFT)은 이론적인 지수적 가속을 제공하지만, 결함 허용(fault-tolerant) 양자 하드웨어의 부재로 인해 여전히 실용적이지 못하다. 따라서 저계수(low-rank) 텐서 표현을 활용하여 이러한 메모리 및 계산 병목 현상을 극복하기 위한 "양자 영감형(quantum-inspired)" 클래식 알고리즘이 필요하다.

방법론
본 논문은 고차 텐서를 표현하기 위해 텐서 트레인(TT) 및 텐서 트레인 연산자(TTO) 형식을 사용하는 SFFT(Superfast Fourier Transform) 기반 균질화 알고리즘의 성능을 조사한다. 연구는 두 단계로 진행된다:

1. 기초 연산 벤치마킹: 저자들은 JAX 프레임워크를 사용하여 Dual Intel Xeon Gold 6240R CPU, NVIDIA A100 GPU, Google TPU v4-8의 세 가지 하드웨어 아키텍처에서 기초적인 TT 대수 연산(덧셈, 곱셈, 수축, 직교화, 압축)을 구현하였다. 두 가지 구현 모드인 "리스트 형식"(코어를 배열의 리스트로 저장)과 "배치 형식"(코어를 단일 배치 배열 내에 저장)을 비교하였다. 정확도를 보장하기 위해 complex64 정밀도를 사용하였으며, 이 과정에서 TPU를 일반적인 BF16 최적화 영역 밖에서 운용하였다. 성능은 실행 시간과 Roofline 모델을 통해 메모리 제한(memory-bound) 또는 연산 제한(compute-bound) 영역을 결정하기 위해 분석되었다.

2. 가속화된 균질화 적용: SFFT 기반 균질화 워크플로우를 이러한 가속기에 적응시켰다. 텐서 계수(rank)가 동적으로 변할 때 발생하는 JAX의 JIT 컴파일 오버헤드를 해결하기 위해 "거친 입도화(coarse-graining)" 전략을 도입하였다. 이는 텐서 계수를 기본 계수($r_0 = 16$)의 배수로 제한하여 재컴파일 이벤트를 최소화하는 방식이다. TPU 구현의 경우, complex64 산술 체계에서 SVD가 높은 이산화 수준에서 수렴하지 않는 현상이 관찰됨에 따라, 표준 SVD 기반 압축을 극성 분해(Polar decomposition) 기반 압축으로 대체하여 수치적 안정성을 확보하였다.

주요 기여

• 최초의 체계적인 TPU 벤치마킹: 본 논문은 GPU 및 CPU와의 직접적인 성능 비교를 포함하여, TPU 하드웨어에서 기초적인 TT 연산을 수행한 첫 번째 엄격한 벤치마킹을 제공한다.
• 하드웨어 가속 TT 대수: 현대적인 가속기 상에서 효율적인 TT 대수 구현을 제시하며, 리스트 형식과 배치 형식 저장 방식의 타당성을 평가하고 특정 성능 특성(예: 메모리 제한 vs 연산 제한 동작)을 식별하였다.
• SFFT 균질화의 실용적 구현: 저자들은 SFFT 기반 균질화 알고리즘을 GPU 및 TPU 실행에 성공적으로 적응시켜, 표준 GPU 기반 FFT 참조 구현으로는 불가능한 규모인 3억 개에서 700억 개의 그리드 포인트를 가진 데이터셋 시뮬레이션을 가능하게 하였다.
• 안정성 분석: TPU 기반 SVD 연산의 complex64 정밀도 하에서의 수치적 불안정성을 식여하고, 고이산화 영역에서 안정적인 대안으로서 극성 분해를 제안하였다.

결과

• 연산 성능:
  • 병렬 연산: 고도로 병렬화 가능한 연산(덧셈, 곱셈, TT-TTO 수축)의 경우, GPU가 높은 이산화 수준에서 우수한 확장성을 보이며 결국 TPU를 추월하였다. TPU는 낮은 이산화 수준에서 낮은 오버헤드를 보였으나, 테스트된 전 범위에서 엄격하게 메모리 제한적(memory-bound)이었다.
  • 직렬 연산: 직렬 연산(직교화, 압축)의 경우, 일반적으로 TPU가 전체 영역에서 GPU보다 우수한 성능을 보였다. 그러나 TPU에서의 SVD 기반 압축은 complex64 정밀도에서 이산화 수준이 $2^7$에 도달할 때 수렴에 실패하였으며, 이는 극성 분대로의 전환을 필연적으로 만들었다.
  • Roofline 분석: GPU는 복잡한 연산에서 주로 연산 제한적(compute-bound)이었던 반면, TPU는 병렬 작업에서는 메모리 제한적이었으나 직렬 작업에서는 큰 이산화 수준에서 연산 제한적 동작으로 전이되었다.
• 균질화 스케일링:
  • GPU 기반 양자 영감형 솔버는 약 700억 개의 그리드 포인트($2^{18}$ 포인트/차원)까지 성공적으로 확장되었으며, 이는 cuFFT 기반 참조 구현의 메모리 한계($2^{12}$ 포인트/차원)를 크게 상회한다.
  • CPU 및 TPU 버전은 각각 메모리 용량 제한으로 인해 $2^{14}$ 및 $2^{10}$ 포인트/차원에 도달하였다.
  • SFFT 방식의 절대적인 실행 시간이 고도로 튜닝된 cuFFT 라이브러리에 비해 아직 완전히 최적화되지는 않았으나, 스케일링 동작은 TT 계수가 적절히 유지되는 분리 가능한 구조의 기하학적 구조에 대해 SFFT 접근 방식이 문제 크기가 커짐에 따라 결국 FFT 기반 방식보다 우수해질 것임을 시사한다.
• 정확도: 압축 컷오프 파라미터에 의해 제어되는 유효 재료 물성에 대해 5% 미만의 상대 오차를 유지하였다.

의의 및 주장
본 논문은 현대적인 가속기 상에서 고성능, 대규모 텐서 기반 균질화를 위한 토대를 마련했다고 주장한다. 텐서 트레인 기법이 산업 규모 시뮬레이션의 메모리 및 계산 병목 현상을 모두 극복할 수 있음을 입증하였으며, 기존 가속기로는 불가능했던 거대 데이터셋의 균질화를 가능하게 한다.

저자들은 본 연구가 근본적인 SFFT 알고리즘을 수정하는 것이 아니라, 효율적인 구현과 가속화에 집중하고 있음을 강조한다. 이들은 이 방법론을 신경 연산자(neural operators) 학습을 위한 정확한 참조 솔루션을 생성할 수 있는 데이터 기반 다중 스케일 모델링의 보완적 도구로 포지셔닝한다. 연구는 현재 이 접근 방식이 저계수 기하학(예: 적층 복합재나 격자 구조의 픽셀화된 미세구조)에 국한되어 있으나, 다중 스케일 재료 모델링을 위한 확장 가능한 물리 기반 양자 영감형 솔버를 향한 유망한 경로를 나타낸다고 결론짓는다. 저자들은 TPU 상의 수치적 안정성 문제를 해결하고 고차 텐서 네트워크로 방법을 확장하기 위한 향-후 연구가 필요함을 언급하며, 임의의 미세구조에 대한 즉각적인 산업적 적용 가능성에 대해서는 신중한 태도를 유지하였다.