🤖 AI

FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU

FlashSinkhorn 는 FlashAttention 스타일의 퓨전과 타이링을 활용하여 HBM 메모리 트래픽을 극적으로 줄여 주며, 최신 베이스라인 대비 최대 161 배의 속도 향상을 달성함과 동시에 대규모 포인트 클라우드 작업을 위한 확장 가능한 최적화를 가능하게 하는 엔트로피 최적 수송을 위한 IO 인지형 GPU 솔버입니다.

원저자: Felix X. -F. Ye, Xingjie Li, An Yu, Ming-Ching Chang, Linsong Chu, Davis Wertheimer

게시일 2026-05-22

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Felix X. -F. Ye, Xingjie Li, An Yu, Ming-Ching Chang, Linsong Chu, Davis Wertheimer

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

두 거대한 인파를 매칭한다고 상상해 보세요. 한 인파는 한쪽 필드 ("소스") 에 서 있고, 다른 인파는 반대쪽 ("타겟") 에 서 있습니다. 당신의 목표는 모든 사람이 이동해야 하는 총 거리를 최소화하도록 모든 사람을 가장 효율적으로 짝지우는 방법을 찾는 것입니다. 이는 최적 수송 (Optimal Transport) 이라는 고전적인 수학 문제입니다.

현대 기계 학습에서는 이 매칭 과정을 조금 더 "부드럽게" 만들어 수학적 처리를 용이하게 하곤 합니다. 이를 엔트로피 최적 수송 (Entropic Optimal Transport) 이라고 합니다. 이를 해결하기 위해 컴퓨터는 싱크혼 반복 (Sinkhorn iterations) 이라는 방법을 사용하는데, 이는 두 인파 사이에서 컴퓨터가 메모를 오가며 주고받는 "뜨거운 감자" 게임과 같습니다. 컴퓨터는 최선의 해법을 찾을 때까지 매칭을 반복적으로 정제합니다.

문제: 교통 체증

이 논문은 이 방법이 소규모 인파에는 잘 작동하지만, 인파가 거대해지면 (수만 명 규모) 거대한 장벽에 부딪힌다고 설명합니다.

컴퓨터의 메모리를 한 도시로 생각해보세요:

HBM(고대역폭 메모리): 이 도시의 주요 고속도로입니다. 규모가 크고 많은 데이터를 보유할 수 있지만, 접근하는 데는 느립니다.
SRAM(온칩 메모리): 컴퓨터 프로세서 내부에 있는 작고 초고속의 개인 사무실입니다. 속도는 매우 빠르지만 용량은 매우 작습니다.

이 매칭 문제를 해결하던 기존 방법들은 매번 한 쌍의 사람을 확인해야 할 때마다 고속도로 (HBM) 에서 사무실 (SRAM) 로, 그리고 다시 고속도로로 이동해야 하는 배송 트럭과 같았습니다. 가능한 쌍이 수백만 개나 되기에 트럭은 고속도로에서 교통 체증에 갇혀 데이터를 끊임없이 왕복시켰습니다. 컴퓨터는 실제 계산을 하는 시간보다 데이터를 기다리는 시간에 더 많은 시간을 보냈습니다.

해결책: FlashSinkhorn

저자들은 FlashSinkhorn 이라는 새로운 도구를 개발했습니다. 그들은 이 매칭 문제의 수학적 구조가 AI 챗봇 (지금 대화 중인 것과 같은) 의 기반 기술인 트랜스포머 (Transformers) 에서 사용되는 수학과 정확히 동일하다는 점을 깨달았습니다.

트랜스포머에는 FlashAttention 이라는 유사한 교통 체증을 해결하는 영리한 트릭이 있습니다. 트럭을 왕복시키는 대신, FlashAttention 은 데이터 전체 "타일 (작은 배치)"을 빠른 사무실로 로드한 뒤, 필요한 모든 계산을 그곳에서 수행하고 최종 결과물만 고속도로에 다시 기록합니다.

FlashSinkhorn 은 바로 이 "타일 기반" 전략을 매칭 문제에 적용합니다:

완전 지도 제거: 모든 가능한 연결을 기록하는 전체 지도 (메모리에 담기엔 너무 큼) 를 작성하는 대신, 한 번에 작은 타일 단위로 연결을 실시간으로 계산합니다.
"사무실" 전략: 현재 계산 배치의 데이터를 빠르고 작은 사무실 (SRAM) 에 유지합니다. 거대한 중간 목록을 느린 고속도로에 기록할 필요 없이, 바로 그곳에서 "매칭 점수"를 업데이트합니다.
스트리밍: 데이터를 컨베이어 벨트처럼 스트리밍하며, 무거운 작업을 처리하고 폐기하는 과정을 진행하면서 고속도로를 비워둡니다.

결과: 속도와 규모

이 논문은 강력한 GPU(특히 A100) 에서 이를 테스트했습니다. 결과는 극적이었습니다:

속도: 초기 계산에서는 기존 최상의 온라인 방법 대비 최대 32 배, 전체 과정 (실수로부터 학습 포함) 에서는 최대 161 배까지 빨랐습니다.
메모리: 기존 방법들은 3 만 명의 인파를 매칭하려다 메모리 부족으로 충돌 (크래시) 했지만, FlashSinkhorn 은 전체 지도를 한 번에 저장하려 하지 않았기에 5 만 명의 인파도 쉽게 처리할 수 있었습니다.
실제 활용: 그들은 수천 장의 이미지와 같은 거대한 데이터셋 비교나 데이터 순서가 뒤섞인 복잡한 회귀 문제 해결과 같은 실제 작업에서도 작동함을 보여주었습니다.

결론

FlashSinkhorn 은 교통 체증에 갇힌 배송 트럭에서 고속 드론으로 업그레이드한 것과 같습니다. 목적지 (수학적 답변은 여전히 정확함) 는 바꾸지 않지만, 데이터를 이동하는 방식 을 바꿉니다. 무거운 작업을 컴퓨터의 빠른 "사무실" 내부에 유지하고 느린 "고속도로"는 최종 결과물에만 사용함으로써, 거대한 매칭 문제를 실용적이고 빠르게 해결합니다. 이로 인해 과거에는 몇 시간이 걸리거나 컴퓨터를 충돌시켰던 작업이 몇 초 만에 완료되도록 바뀐 것입니다.

기술 요약: FlashSinkhorn: GPU 기반 IO 인지 엔트로피 최적 수송

문제 제기

Sinkhorn 반복을 통해 해결되는 엔트로피 최적 수송 (EOT) 은 생성 모델링부터 분포 정렬에 이르기까지 확률 측도를 비교하는 현대 기계 학습의 핵심 도구입니다. 그러나 GPU 에서 대규모 점 구름 (예: $n \approx 70,000$ ) 으로 EOT 를 확장하는 것은 지속적인 병목 현상으로 남아 있습니다.

기존 GPU 솔버는 두 가지 주요 비효율성에 직면해 있습니다:

텐서화 구현 (Tensorized Implementations): 이들은 밀집된 $n \times m$ 비용 행렬을 물리적으로 생성하여 이차적인 대역폭 메모리 (HBM) 트래픽 ($O(nm) $읽기/쓰기) 과 메모리 용량 문제를 초래하며,$ n$이 증가함에 따라 메모리 부족 (OOM) 오류를 유발합니다.
온라인 백엔드 (Online Backends): KeOps 나 OTT-JAX 를 통해 상호작용을 실시간으로 평가하여 밀집 행렬을 저장하지는 않지만, 범용 타일링 맵-리듀스 축소 커널에 의존합니다. 이러한 커널은 제한된 퓨전 (fusion) 으로 인해 텐서 코어 GEMM 구조를 활용하지 못해 계산 활용도가 낮고 메모리 사용량이 감소했음에도 실행 시간이 길어지는 결과를 초래합니다.

핵심 시스템 병목 현상은 현재 방법론들이 과도한 HBM 트래픽으로 인해 메모리 바운드되거나, 낮은 명령어 효율성으로 인해 계산 바운드되어 하류 파이프라인에서 필요한 반복적인 대규모 해결 작업에 확장하지 못한다는 점입니다.

방법론

저자들은 제곱 유클리드 비용을 가진 EOT 를 위한 IO 인지 GPU 솔버인 FlashSinkhorn을 소개합니다. 이 방법론은 안정화된 로그 도메인 Sinkhorn 업데이트와 Transformer 의 어텐션 메커니즘 사이의 직접적인 유사성을 도출하여 FlashAttention의 원리를 활용하여 메모리 접근 패턴을 최적화합니다.

1. 어텐션 형태 재형식화

핵심 이론적 통찰 (명제 1) 은 제곱 유클리드 비용의 경우, 안정화된 로그 도메인 Sinkhorn 업데이트가 편향된 내적 점수 행렬의 행별 LogSumExp (LSE) 축소로 재작성될 수 있다는 것입니다.

구체적으로, 업데이트 $f_i \leftarrow -\varepsilon \text{LSE}_j [(g_j - C_{ij})/\varepsilon + \log b_j]$ 는 $Q=\sqrt{2}X$ , $K=\sqrt{2}Y$ 이며 $\delta, \gamma$ 는 사전 계산된 로그 가중 편향인 행렬 $S_X = (QK^\top + \mathbf{1}(\hat{g} + \delta))/\varepsilon$ 의 LSE 를 계산하는 것과 동일합니다.
이 재형식화는 Sinkhorn 반복을 확장된 내적 어텐션에 사용되는 정확한 정규화 통계량으로 변환하여 IO 인지 스트리밍 기법의 적용을 가능하게 합니다.

2. 스트리밍 커널 및 타일링

FlashSinkhorn 은 느린 HBM 과 빠른 온칩 SRAM 을 사용하는 2 단계 메모리 모델을 채택한 퓨전 Triton커널을 사용하여 이 재형식화를 구현합니다:

물리화 방지: $n \times m$ 비용 행렬은 HBM 에 결코 물리화되지 않습니다. 대신 $Q$ 와 $K$ 의 타일들이 스트리밍됩니다.
온라인 LSE: 커널은 점수 행렬의 타일을 처리하는 동안 온칩 (SRAM) 에서 실행되는 행별 최댓값 및 합 지수 통계를 유지합니다. 이를 통해 HBM 접근 횟수를 크게 줄이면서 정확한 LSE 축소를 계산할 수 있습니다.
퓨전: 전체 업데이트 (점수 계산, 편향 추가, LSE 축소, 잠재적 업데이트) 는 단일 커널 패스로 퓨전됩니다. 이는 HBM 에 대한 중간 쓰기를 제거하고 커널 시작 오버헤드를 줄입니다.
행 고정 탐색 (Row-Stationary Traversal): FlashAttention 이 종종 바깥쪽 루프에서 키/값을 반복하는 것과 달리, FlashSinkhorn 은 Sinkhorn 업데이트의 행별 축소 특성에 더 잘 부합하도록 행 고정 탐색 (SRAM 에 쿼리 행 블록을 고정하고 모든 키 블록을 스트리밍) 을 채택합니다.

3. 미분 및 2 차 최적화

이 논문은 이 IO 인지 전략을 미분 및 2 차 방법론으로 확장합니다:

수송 행렬 적용: 바리센트릭 투영과 그래디언트 계산은 어텐션과 유사한 연산 (가중 평균) 으로 표현되어, 수송 행렬 $P^*$ 를 물리화하지 않고도 동일한 스트리밍 커널을 통해 계산될 수 있습니다.
헤시안 - 벡터 곱 (HVP): 저자들은 헤시안 - 벡터 곱을 수송 - 벡터 및 수송 - 행렬 곱의 시퀀스로 분해하는 스트리밍 HVP 오라클을 유도합니다. 이는 $O(n^2 d^2)$ 크기의 헤시안 텐서나 $O(nm) $크기의 수송 행렬을 물리화하지 않고 작업 메모리를$ O((n+m)d)$로 줄입니다.

주요 기여

어텐션 형태 Sinkhorn 업데이트: 저자들은 안정화된 Sinkhorn 업데이트를 편향된 내적 LSE 축소로 표현하여 EOT 를 위한 정확한 FlashAttention 스타일 스트리밍을 가능하게 합니다.
IO 인지 EOT 솔버: 스트리밍된 Sinkhorn 반복을 위한 퓨전 Triton 커널과 미분을 위한 수송 행렬 없는 스트리밍 연산자를 제공하여 HBM IO 볼륨을 대폭 줄입니다.
확장 가능한 2 차 최적화: 스트리밍 HVP 를 도입하여 EOT 기반 손실 함수에 대한 확장 가능한 Newton-CG 및 신뢰 영역 방법을 가능하게 합니다. 이는 이전에는 메모리 제약으로 인해 대규모에서 실현 불가능했습니다.
대규모 평가: 합성 벤치마크와 OTDD(Optimal Transport Dataset Distance) 및 셔플 회귀와 같은 하류 작업 전반에 걸쳐 A100 GPU 에서 상당한 속도 향상과 메모리 효율성을 입증했습니다.

실험 결과

NVIDIA A100-80GB GPU 에서 평가된 FlashSinkhorn 은 최신 기법 (GeomLoss 의 KeOps/텐서화 백엔드 및 OTT-JAX) 대비 상당한 성능 향상을 보여줍니다:

속도:
- 순방향 통과: KeOps 대비 최대 32 배, OTT-JAX 대비 5.1 배의 속도 향상을 달성합니다.
- 종단 간 (순방향 + 역방향): 고차원 ( $d=512$ ) 에서 KeOps 대비 최대 161 배의 속도 향상을 달성합니다.
- HVP: 차원에 따라 3~25 배의 속도 향상을 달성하여, 다른 방법들이 시간 초과가 발생하는 규모에서도 Newton-CG 를 실용적으로 만듭니다.
메모리 효율성:
- FlashSinkhorn 은 선형 메모리 복잡도 $O((n+m)d)$ 를 유지하는 반면, 텐서화 기법들은 $O(nm) $또는$ O(n^2)$로 확장됩니다.
- 이는 텐서화 방법이 메모리 부족을 겪고 온라인 기법 (KeOps) 이 시간 제한을 초과하는 $n=50,000$ 개의 점 구름과 최대 $d=1024$ 의 차원을 성공적으로 처리합니다.
하류 작업:
- OTDD: FlashSinkhorn 은 피크 메모리가 1GB 미만인 상태에서 OTDD 거리 계산 및 그래디언트 흐름을 $n=60,000$ (MNIST $\leftrightarrow$ Fashion-MNIST) 으로 확장하는 반면, 텐서화 기법들은 $n > 20,000$ 에서 메모리 부족 (OOM) 을 겪습니다.
- 안장점 탈출: 셔플된 선형 회귀에서 스트리밍 HVP 는 안장점을 감지하기 위해 최소 헤시안 고유값을 모니터링할 수 있게 하여, 하이브리드 Adam-Newton 최적화기가 안장점을 탈출하고 빠르게 수렴하도록 합니다.

중요성 및 주장

이 논문은 FlashSinkhorn 이 메모리 바운드 영역에서 대규모 EOT 를 실용화한다고 주장합니다. Sinkhorn 업데이트를 FlashAttention 의 IO 인지 구조에 맞게 재작성함으로써, 저자들은 최적 수송과 현대적인 트랜스포머 가속 기법 간의 격차를 해소했습니다.

그 중요성은 다음과 같습니다:

시스템 수준 최적화: 이러한 이득은 Sinkhorn 반복 자체의 알고리즘적 변경이 아닌 커널 수준의 전문화 (퓨전 및 타일링) 에 기인합니다.
새로운 워크로드 가능: 메모리 발자국 및 실행 시간의 감소는 데이터셋 거리 및 암시적 미분과 같은 하류 작업에서 EOT 를 이전에 계산적으로 금지적이었던 규모에서 사용할 수 있게 합니다.
재현성: 저자들은 확장 가능한 최적 수송에 대한 추가 연구를 촉진하기 위해 오픈 소스 구현을 공개합니다.

저자들은 겸손하게도 이 접근 방식이 현재는 특정 구조를 가진 비용 (정규화 후 제곱 유클리드 및 코사인 거리) 에만 적용된다고 지적합니다. 이러한 구조를 갖지 않는 비용 (예: 제곱근이 포함된 원시 유클리드 거리 또는 학습된 신경 비용) 은 향후 과제로 식별되었습니다.