Implicit differentiation of tensor network algorithms

기술 요약: 텐서 네트워크 알고리즘의 암시적 미분 (Implicit Differentiation)

문제 정의
바닥 상태 계산을 위한 투영된 얽힘 쌍 상태(Projected Entangled-Pair States, PEPS)의 변분 최적화는 국소적 변분 파라미터에 대한 에너지 기울기(gradient)를 정확하게 계산하는 것에 크게 의존한다. 현재의 최첨단 방식은 이 기울기를 계산하기 위해 자동 미분(Automatic Differentiation, AD)을 활용한다. 그러나 이 방법은 세 가지 중요한 문제에 직면해 있다:

1. 효율성: 표준 역방향 모드(reverse-mode) AD는 모든 수축 알고리즘(예: CTMRG 또는 Boundary MPS)의 반복 단계마다 역전파(backpropagation)를 수행해야 한다. 고정점(fixed-point) 정식화를 통해 전체 계산 그래프를 저장하는 것을 피하더라도, 희소 고유값 솔버(sparse eigensolver)나 SVD와 같은 반복적인 서브루틴의 존재는 기울기 계산 과정에서 중첩된 반복 절차를 생성하여 높은 계산 비용을 초득한다.
2. 안정성: 수축 알고리즘에 사용되는 원시 서브루틴(예: 절단된 고유값 분해 또는 SVD)을 미분하는 것은 수치적 불안정성을 유발한다. 이러한 불안정성은 종종 근접한 퇴화된(near-degenerate) 고유값이나 특이값으로부터 발생하며, 여기서 서브루틴의 미분값은 발산한다. 결정적으로, 이러한 불안정성은 미분 절차에 의한 인위적인 산물이다. 물리적인 에너지 기울기 자체는 기저의 변화에 대해 안정적이다.
3. 구현 복잡성: PEPS 최적화에 AD를 적용하려면 수축 알고리즘의 모든 원시 연산이 AD 엔진과 호환되어야 한다. 네이티브가 아닌 서브루틴에 대해 커스텀 미분 규칙을 정의하는 것은 매우 까다로우며 유연성을 제한한다.

방법론
저자들은 이러한 과제들을 해결하기 위해 **암시적 미분(implicit differentiation)**에 기반한 프레임워크를 제안한다. 이 방법은 수축 알고리즘의 반복적인 단계들을 미분하는 대신, 수렴된 수축 환경을 단일 대수적 특성 방정식(algebraic characteristic equations)의 해로 재정의한다.

핵-단계는 다음과 같다:

1. 대수적 특성화: 수렴된 환경(예: CTMRмG의 코너 및 엣지 텐서, 또는 경계 MPS 텐서)은 반복적인 업데이트 규칙 $x = f(x, p)$가 아니라, $F(y, p) = 0$이라는 일련의 방정식 집합으로 특성화된다. 여기서 $y$는 변분 파라미터 $p$에 암시적으로 의존하는 환경 변수들(중간 변수인 등거리 사상(isometry)이나 특이값 포함)을 나타낸다.
2. 암시적 미분: $F(y, p) = 0$에 암시적 함수 정리(implicit function theorem)를 적용함으로써, 기울기 계산은 단일한, 중첩되지 않은 선형 방정식 시스템을 푸는 것으로 축소된다:
   $$\tilde{y} \partial_y F = -\tilde{y} \partial_p F$$
   여기서 $\tilde{y}$는 수반 변수(adjoint variable)이다. 이를 통해 반복적인 서브루틴을 통한 역전파를 제거할 수 있다.
3. 안정적 파라미터화: 핵심 혁신은 변수 집합 $y$의 선택과 특성 방정식의 파라미터화에 있다. 저자들은 수축 환경의 내재적인 게이지 대칭성(gauge symmetries)을 활용하여, 수치적 불안정성의 원인이 되는 자유도(예: 퇴화된 부분 공간 내의 변화)를 명시적으로 제거하는 파라미터화를 구축한다. 예를 들어, CTMRG에서 등거리 사상의 변화는 절단된 고유값 솔버의 미분과 관련된 발산 항을 우회하도록 그 널 공간(null space) 내로 제한된다.
4. 모듈형 구현: 이 접근 방식은 수축 알고리즘과 기울기 계산을 분리한다. 순방향 패스(forward pass)는 변경되지 않은 채 기존의 수축 루틴을 그대로 사용하며, 역방향 패스(reverse pass)는 대수적 특성 방정식(단순 텐서 수축으로 구성됨)만을 미분하면 된다.

주요 기여
본 논문은 세 가지 주요 PEPS 수축 스킴에 대한 특성 방정식의 명시적 정식화를 제공한다:

• $C_{4v}$ 대칭 CTMRG: 저자들은 코너 텐서, 엣지 텐서, 그리고 파라미터화된 등거리 사상을 포함하는 특성 방정식 집합을 유도한다. 이 정식화는 절단된 고유값 분해(eigh)를 미분할 필요를 없애고, 퇴화된 고유값과 관련된 불안정성을 해결한다.
• 경계 MPS (VUMPS): 상단 및 하단 경계 MPS 고유벡터와 혼합 환경에 대한 특성 방정식 집합이 정식화되었다. 이 접근 방식은 기울기 계산 중에 발생하는 경계 MPS 텐서의 게이지 고정(gauge-fixing) 문제를 피하고, 지배적인 고유값 솔버로부터 발생하는 중첩된 선형 문제를 제거한다.
• 비대칭 CTMRG (대칭성 없음): 이 방법은 공간적 대칭성이 없는 일반적인 시스템으로 확장된다. 저자들은 게이지 불변성과 절단된 SVD를 미분할 때의 안정성을 보장하기 위해 수정된 코너 및 엣지 텐서와 커스-스퀘어 루트(custom "square root") 변수를 도입한다.

결과
저자들은 PEPSKit.jl 및 MPSKit.jl 패키지를 사용하여 암시적 미분 방식을 고정점 미분 및 직접적인 역방향 모드 AD와 벤치마킹하였다:

• 정확도: 암시적 미분으로 계산된 기울기는 선형 솔버의 허용 오차 범위 내에서 참조 수치 미분(finite-difference) 값 및 블랙박스 역방향 모드 AD와 일치하였으며, 이는 중간 정도 수렴된 환경에서도 마찬가지였다.
• 성능:
  • $C_{4v}$ CTMRG 및 경계 MPS의 경우, 암시적 방식이 고정점 미분보다 일관되게 우수한 성능을 보였으며, 결합 차원($D$)과 환경 차원($\chi$)이 커질수록 속도 향상이 두드러졌다.
  • 비자명한 단위 셀(unit cell)에 대한 비대칭 CTMRG의 경우, 암시적 방법은 우수한 점근적 스케일링을 보여준다. 고정점 미분은 구현 오버헤드로 인해 낮은 차원에서 더 빠를 수 있지만, 희소 행렬 분해가 필요한 높은 차원에서는 암시적 방법이 훨씬 더 효율적이 된다.
• 스케일링: 많은 경우, 암시적 미분을 사용한 기울기 계산 비용은 수축 비용 자체에 비해 부차적인 단계가 되지만, 고정점 미분은 종종 지배적인 병목 현상으로 남는다.

의의 및 주장
본 논문은 이 프레임워크가 텐서 네트워크 프로그래밍의 패러다임 전환을 의미한다고 주장한다. 알고리즘적 서브루틴이 아닌 대수적 최적 조건(optimality conditions)을 미분함으로써, 이 방법은 다음을 달 수 있다:

• 구현 단순화: 수축 알고리즘의 모든 서브루틴이 미분 가능할 필요가 없으므로, 순방향 패스에서 표준적이고 고도로 최적화된 선형 대수 루틴을 사용할 수 있게 한다.
• 안정성 향상: 게이지 대칭성을 활용하여 안정적인 파라미터화를 구축함으로써, 서브루틴의 미분에서 발생하는 수치적 불안정성을 자연스럽게 제거한다.
• 효율성 개선: 선형 문제를 비중첩화(un-nesting)함으로써 기울기 평가의 계산 복잡도를 줄이고, 시스템 크기에 따른 스케일링을 개선한다.

저자들은 이 접근 방식이 세 가지 특정 수축 스킴에 대해 시연되었지만, 핵심 알고리즘(Algorithm 1 및 2)은 미분 가능한 최적 조건에 의해 해가 정의될 수 있는 모든 최적화 문제에 적용 가능한 일반적인 것이라고 강조한다. 이들은 모든 PEPS 문제를 해결했다고 주장하는 것이 아니라, 이 접근 방식이 안정적이고 기울기 기반의 PEPS 최적화를 훨씬 더 실용적이고 견고하게 만든다고 주장한다.