[논문 리뷰] Importance Sparsification for Sinkhorn Algorithm

Mengyu Li, Jun Yu외 2명|arXiv (Cornell University)|2023. 06. 11.

Bayesian Modeling and Causal Inference인용 수 5

한 줄 요약

이 논문은 엔트로피 정규화된 최적 운반 문제(OT) 및 비균형 OT(UOT)의 Sinkhorn 알고리즘을 가속화하기 위해, 샘플링 확률을 위한 천연 상한을 활용해 희소 커널 행렬을 구성하는 새로운 중요도 희소화 방법인 Spar-Sink을 제안한다. 이 방법은 반복당 계산 비용을 $O(n^2)$에서 $\tildewide{O}(n)$으로 감소시키며, 온건한 조건 하에서 일致한 추정기를 제공하고, 합성 및 실제 에코카디오그래프 데이터(심장 사이클 분석 및 생성 모델링 포함)에서 경쟁 기법들보다 빠르고 정확한 성능을 보인다.

ABSTRACT

Sinkhorn algorithm has been used pervasively to approximate the solution to optimal transport (OT) and unbalanced optimal transport (UOT) problems. However, its practical application is limited due to the high computational complexity. To alleviate the computational burden, we propose a novel importance sparsification method, called Spar-Sink, to efficiently approximate entropy-regularized OT and UOT solutions. Specifically, our method employs natural upper bounds for unknown optimal transport plans to establish effective sampling probabilities, and constructs a sparse kernel matrix to accelerate Sinkhorn iterations, reducing the computational cost of each iteration from $O(n^2)$ to $\widetilde{O}(n)$ for a sample of size $n$. Theoretically, we show the proposed estimators for the regularized OT and UOT problems are consistent under mild regularity conditions. Experiments on various synthetic data demonstrate Spar-Sink outperforms mainstream competitors in terms of both estimation error and speed. A real-world echocardiogram data analysis shows Spar-Sink can effectively estimate and visualize cardiac cycles, from which one can identify heart failure and arrhythmia. To evaluate the numerical accuracy of cardiac cycle prediction, we consider the task of predicting the end-systole time point using the end-diastole one. Results show Spar-Sink performs as well as the classical Sinkhorn algorithm, requiring significantly less computational time.

연구 동기 및 목표

• 대규모 최적 운반 문제(OT) 및 비균형 OT(UOT)에서 Sinkhorn 알고리즘의 높은 계산 비용을 해결하기 위해.
• 정확도를 유지하면서 계산 복잡도를 극적으로 감소시키는 희소 근사 방법을 개발하기 위해.
• 에코카디오그래프 분석 및 생성 모델링과 같은 실제 응용 분야에서 효율적이고 정확한 OT 및 UOT 계산을 가능하게 하기 위해.
• 온건한 정규성 조건 하에서 제안된 방법의 이론적 일致성을 확립하기 위해.
• 합성 및 실제 데이터셋에서 기존 최첨단 기법들보다 속도와 추정 오차 측면에서 뛰어난 성능을 내기 위해.

제안 방법

• 미지의 최적 운반 계획에 대한 자연적 상한을 활용해 효과적인 샘플링 확률을 유도함으로써 커널 행렬의 희소화를 위한 Spar-Sink를 제안한다.
• 이러한 중요도 확률 기반으로 원소 수준의 샘플링을 수행하여 희소 커널 행렬을 구성함으로써, Sinkhorn 반복당 계산 비용을 $O(n^2)$에서 $\tildewide{O}(n)$으로 감소시킨다.
• 이 희소화된 커널을 Sinkhorn 알고리즘에 적용하여 엔트로피 정규화된 OT 및 UOT 해를 효율적으로 근사한다.
• 생성 모델링에서는 $s = 10s_0(n)$, 색상 이동에서는 $s = 8s_0(n)$의 서브샘플링 파라미터를 사용하여 정확도와 속도 사이의 균형을 이룬다.
• OT 프레임워크에서 비교를 위해 강건한 변형인 Robust-NysSink을 활용하고, Spar-Sink을 새로운 오토인코더 변형인 Spar-Sink 오토인코더(SSAE)에 통합한다.
• 이론적 분석을 통해 온건한 정규성 조건 하에서 정규화된 OT 및 UOT에 대한 제안된 추정기의 일치성을 입증한다.

실험 결과

연구 질문

• RQ1 천연 상한을 통한 중요도 희소화가 추정 정확도를 희생시키지 않고 Sinkhorn 알고리즘의 계산 비용을 상당히 감소시킬 수 있는가?
• RQ2 Nys-Sink 및 Robust-NysSink과 같은 기존 변종 대비 Spar-Sink이 OT 및 UOT 문제에서 속도와 정확도 측면에서 어떻게 비교되는가?
• RQ3 Spar-Sink은 에코카디오그래프 데이터에서 심장 사이클 분석을 효과적으로 가속화하여, 심장수축기 말기 시점을 심장확장기 말기에서 정확하게 예측할 수 있는가?
• RQ4 Spar-Sink을 생성 모델(SAE)에 통합한 SSAE는 원래의 Sinkhorn 오토인코더(SAE) 대비 성능을 유지하거나 향상시키는가?
• RQ5 온건한 정규성 조건 하에서 Spar-Sink 추정기의 정규화된 OT 및 UOT에 대한 이론적 일치성은 어떠한가?

주요 결과

• 색상 이동 작업에서, 운반 계획 계산에 소요된 CPU 시간을 Sinkhorn의 60.45초에서 Spar-Sink로 3.15초로 감소시켰고, 시각적 결과는 거의 동일했다.
• 에코카디오그래프 데이터에서, Spar-Sink는 전통적인 Sinkhorn과 유사한 정확도로 심장수축기 시간을 예측했지만, 훨씬 더 짧은 시간에 수행했다.
• 생성 모델링에서, SSAE는 FID 점수 23.65 ± 0.06을 기록하여 SAE의 24.72 ± 0.13보다 낮게 유지했고, 학습 시간도 절반 이하인 64.81초로 완료했다.
• 이 방법은 반복당 $\tildewide{O}(n)$의 계산 비용을 달성하여, 기존 Sinkhorn의 $O(n^2)$ 비용에 비해 상당한 향상이 있었다.
• 합성 데이터에 대한 실험을 통해 Spar-Sink가 추정 오차와 계산 속도 측면에서 주류 경쟁 기법들을 모두 능가함을 확인했다.
• 이론적 분석을 통해 Spar-Sink 추정기의 정규화된 OT 및 UOT에 대해 온건한 정규성 조건 하에서 일치성이 입증되었다.