FlashSinkhorn 2: 블록 희소 엔트로피 최적 수송으로 1억 입자를 한 장의 GPU에 올리기

arXiv 번호 2610.02395로 공개된 FlashSinkhorn 2: Block-Sparse Entropic Optimal Transport는 저차원 점구름 위에서 제곱 유클리드 비용을 쓰는 대규모 이산 엔트로피 최적 수송 문제를 한 장의 GPU에서 푸는 방법을 제안합니다. 이 논문이 내세우는 주장은 하나의 문장으로 정리할 수 있는데, 성긴 중심점 단계와 블록 희소 정밀 단계로 계산량을 나누면 모든 쌍을 매번 계산하지 않고도 수송 해를 끝까지 수렴시킬 수 있다는 것입니다. 원문 논문으로 이어지는 링크는 이 글 제목 아래에 함께 붙어 있습니다.

제가 이 논문을 읽기 시작한 이유는 단순했습니다. 최적 수송은 개념은 우아한데 계산량은 늘 야박했기 때문입니다. 분포 두 개를 주면 한쪽 질량을 다른 쪽으로 옮기는 가장 싼 방법을 찾는다는 아이디어 자체는 직관적인데, 막상 입자가 수백만 개를 넘어가면 행렬 하나를 메모리에 올리는 일부터 막힙니다. 그런데 이 논문은 1억 단위를 한 장의 A100에서 풀었다고 말합니다. 이 숫자는 그대로 믿기 전에 조건을 봐야 합니다. 그래서 저는 숫자보다 그 뒤의 구조가 더 중요하다고 봅니다.

왜 매 반복마다 모든 쌍을 봐야 했을까

엔트로피 최적 수송을 처음 접하면 Sinkhorn 반복이 낯설게 느껴질 수 있습니다. 그런데 그 뼈대는 생각보다 단순합니다. 출발 분포의 각 입자가 도착 분포의 각 입자에게 얼마나 질량을 보낼지를 행렬로 두고, 행 합과 열 합이 주어진 주변 분포와 맞을 때까지 행과 열을 번갈아 맞추는 것입니다. 이때 두 점 사이의 비용이 행렬의 각 칸에 들어가는데, 엔트로피 정칙화가 들어가면 커널 행렬이 조밀해지면서 모든 칸이 0이 아닌 값을 갖게 됩니다.

문제는 그다음입니다. 입자가 n개와 m개면 매 반복마다 n 곱하기 m개의 쌍을 봐야 합니다. 입자가 만 개면 1억 칸이고, 백만 개면 1조 칸입니다. 이 계산을 메모리에 올려두면 당연히 터지고, 흘려보내듯 그때그때 계산해도 연산량 자체는 그대로 남습니다. 이전 세대의 스트리밍 해법인 FlashSinkhorn은 바로 이 저장 문제를 풀었습니다. 조밀한 커널을 메모리에 보관하지 않고 타일 단위로 흘리면서 계산했기 때문입니다. 메모리 벽은 낮아졌는데, 연산 벽은 그대로였습니다. 매 반복마다 여전히 모든 쌍을 평가해야 했기 때문입니다.

저는 이 지점에서 오래된 질문을 떠올렸습니다. 정말 모든 쌍이 매번 필요할까요. 멀리 떨어진 점들 사이의 커널 값은 대개 아주 작아서, 굳이 매번 정밀하게 계산하지 않아도 주변 분포를 맞추는 데는 거의 영향을 주지 않기 때문입니다. 그렇다면 누가 멀고 누가 가까운지를 미리 알 수 있다면, 계산을 건너뛰어도 되지 않을까요. FlashSinkhorn 2는 이 질문에서 출발합니다.

커널을 저장하지 않아도 여전히 무거웠던 이유

스트리밍 계산에 익숙하지 않으면 여기서 오해가 생길 수 있습니다. 커널을 저장하지 않으면 문제가 끝난 것 아닌가요. 실제로는 그렇지 않습니다. 저희가 새로운 AI 아키텍처를 볼 때도 benchmark 숫자만 보지는 않습니다. 실제로 무엇이 바뀌었고, 그 변화가 시스템에서 어떤 비용 구조를 만드는지를 함께 봅니다. FlashSinkhorn의 경우 바뀐 것은 메모리 사용량이었고, 바뀌지 않은 것은 반복당 연산량이었습니다.

이 차이는 생각보다 중요합니다. GPU는 메모리가 모자라면 타일로 나눠서 돌리면 되지만, 연산량이 그대로면 시간이 그대로 걸리기 때문입니다. 특히 저차원 점구름에서는 점들이 공간에 퍼져 있어서, 어떤 쌍은 가까이 붙어 있고 어떤 쌍은 아주 멀리 떨어져 있습니다. 그런데 조밀한 반복은 이 둘을 구분하지 않고 매번 같은 정성으로 계산합니다. 가까운 쌍도 보고 먼 쌍도 보는 것이죠. 먼 쌍의 기여가 작다는 것을 알면서도 매번 들여다보는 셈입니다.

여기서 재미있는 질문이 생깁니다. 먼 쌍을 아예 계산하지 않으면 행 합과 열 합이 틀어지지 않을까요. 주변 분포를 맞추는 것이 Sinkhorn의 목표인데, 계산을 생략하면 그 목표가 흔들릴 수 있기 때문입니다. 그래서 생략에는 보증이 필요합니다. 어느 쌍을 건너뛰어도 행과 열의 오차가 어느 수준 이하로 남는다는 보장 말입니다. 이 논문은 그 보장을 블록 단위로 만들었다는 점에서 이전과 다릅니다.

성긴 해상도에서 먼저 풀면 무슨 일이 생길까

FlashSinkhorn 2의 첫 번째 단계는 공간을 셀로 나누고 각 셀의 중심점에서 먼저 푸는 것입니다. 점이 수천만 개라도 셀 개수는 훨씬 적기 때문에, 중심점 사이의 문제는 크기가 작아 빠르게 풀립니다. 그러면 각 셀을 대표하는 퍼텐셜 값이 나오는데, 이를 다시 셀 안의 모든 점으로 펼쳐주는 과정을 거칩니다. 중심점이 배운 가격표를 셀 안의 점들이 그대로 물려받는다고 생각하면 됩니다.

생각해보면 사람도 비슷하게 배웁니다. 새로운 도시의 교통비를 물을 때 모든 골목을 따로 묻지 않고, 먼저 구 단위 시세부터 파악한 뒤에 세부 주소를 조정하기 때문입니다. 계산도 같은 순서로 가면 효율적입니다. 전체를 처음부터 정밀하게 푸는 대신 성긴 해상도에서 윤곽을 잡고 내려오는 것이죠.

그런데 여기서 바로 다음 의문이 생깁니다. 셀 안의 점들이 모두 같은 값을 물려받아도 될까요. 셀이 크면 셀 안에서도 점들의 위치가 꽤 다를 수 있기 때문입니다. 그래서 이 논문은 들어 올린 해를 그대로 믿지 않습니다. 표본을 뽑아 주변 분포 검사를 돌리고, 들어 올린 퍼텐셜이 실제 점 수준의 조건을 만족하는지 확인합니다. 통과하면 그대로 다음으로 가고, 통과하지 못하면 중심점 단계로 돌아가서 더 돌립니다. 저는 이 검사가 논문의 태도를 보여준다고 봅니다. 성긴 해상도를 쓰되 맹신하지 않겠다는 것이죠.

들어 올린 해를 믿어도 될지 어떻게 확인할까

표본 주변 검사가 왜 필요한지 조금 더 들여다볼 필요가 있습니다. 중심점에서 푼 해는 정의상 중심점에 대해서만 정확합니다. 그런데 우리가 진짜 맞춰야 하는 것은 개별 점들의 주변 분포입니다. 중심점 개수가 셀 개수만큼이라면, 점 개수는 그 수백 배나 수천 배가 될 수 있습니다. 이 간극을 메우지 않고 내려오면 행 합과 열 합이 어긋나기 때문입니다.

그래서 연구진은 전체 점을 다 검사하는 대신 일부를 뽑아 검사합니다. 전체를 검사하면 성긴 단계의 이득이 사라지기 때문입니다. 표본으로 추정한 주변 오차가 기준 안에 들어오면 들어 올리기를 받아들이고, 그렇지 않으면 중심점 업데이트를 계속합니다. 이 판정은 매번 전체를 보는 대신 일부만 보고 방향을 정한다는 점에서 통계적인 타협입니다. 표본이 나쁘면 판정이 흔들릴 수 있지만, 전체를 매번 보는 비용을 생각하면 합리적인 선택입니다.

실제 시스템 관점에서 보면 다음 질문이 바로 생깁니다. 이 검사를 통과했다고 해서 진짜 수렴이 보장될까요. 통과는 어디까지나 들어 올린 해가 크게 어긋나지 않았다는 신호일 뿐, 남은 오차가 없다는 뜻은 아니기 때문입니다. 논문도 이 점을 인정하는 구조로 되어 있습니다. 그래서 다음 단계가 필요합니다. 성긴 단계가 잡은 윤곽 위에 남는 오차를 지우는 정밀 단계 말입니다.

성긴 격자가 지우지 못하는 오차는 어디에 남을까

중심점에서 계속 돌리면 오차가 계속 줄 것 같지만, 실제로는 그렇지 않습니다. 중심점 업데이트는 중심점 해상도에서 보이는 오차만 고칠 수 있기 때문입니다. 셀 안의 점들이 중심점과 조금씩 어긋나는 부분, 특히 셀 경계 근처에서 갈라지는 부분은 중심점 해상도에서는 보이지 않습니다. 아무리 중심점에서 반복해도 손이 닿지 않는 오차인 셈입니다.

표면적으로 보면 단순합니다. 실제로는 그렇지 않습니다. 이 남은 오차는 전체 점 수준에서 보면 주변 분포의 어긋남으로 나타나고, 목표 잔차까지 가려면 결국 점 수준의 계산이 필요하기 때문입니다. 그런데 점 수준으로 돌아가면 다시 n 곱하기 m의 벽 앞에 서게 됩니다. 성긴 단계로 아낀 시간을 정밀 단계에서 다 써버리면 남는 것이 없기 때문입니다.

제가 이 논문에서 가장 흥미롭게 본 부분은 바로 여기입니다. 성긴 단계를 빠르게 돌리는 것 자체보다, 성긴 단계가 닿지 못하는 오차를 어떻게 싸게 지우는지가 이 설계의 승부처이기 때문입니다. 논문은 여기서 블록 희소 단계로 넘어갑니다. 모든 점 쌍을 보는 대신, 꼭 봐야 하는 블록만 골라서 보는 것입니다.

공간을 어떻게 자르면 건너뛰어도 안전할까

블록 희소 단계의 출발점은 공간 순서입니다. 점들을 Morton 순서로 정렬하면 가까이 있는 점들이 메모리에서도 가까이 모입니다. Z자 모양으로 공간을 훑으면서 번호를 매긴다고 생각하면 됩니다. 그러면 행렬의 블록 하나가 공간상의 가까운 묶음과 묶음 사이의 상호작용을 나타내게 됩니다. 가까운 묶음은 값이 크고, 먼 묶음은 값이 작아지는 구조가 만들어지는 것이죠.

이 정렬이 끝나면 스크리닝이 가능합니다. 블록마다 기여도를 어림잡아서, 작을 것으로 예상되는 블록은 아예 계산하지 않는 것입니다. 가까운 블록만 남기고 먼 블록을 비우면, 조밀했던 행렬이 듬성듬성한 블록 희소 행렬로 바뀝니다. 계산량은 남긴 블록 개수에 비례해서 줄어들기 때문에, 먼 상호작용이 많은 저차원 문제에서는 절감이 큽니다.

여기서는 조금 조심해서 읽을 필요가 있습니다. 가까운 것만 남기면 정확도가 떨어지는 것 아닌가요. 직관적으로는 그렇습니다. 그래서 무엇으로 남길지를 정하는 기준이 중요합니다. 단순히 거리만 보면 질량이 몰린 곳을 놓칠 수 있기 때문입니다. 질량이 큰 블록은 거리가 좀 있어도 기여가 클 수 있어서, 거리와 질량을 함께 봐야 합니다. 논문이 블록 질량으로 임계값을 정했다는 대목이 바로 이 때문입니다. 얼마나 먼지뿐 아니라 얼마나 무거운지를 함께 따진다는 것이죠.

생략해도 된다는 것을 어떻게 보장할까

이 논문에서 제가 기술적으로 가장 눈여겨본 대목은 생략된 타일의 기여도를 행과 열마다 묶어두었다는 점입니다. 블록 질량으로 정한 임계값이 각 행과 각 열에서 생략된 타일이 만들 수 있는 오차의 상한이 되기 때문입니다. 한두 개의 행에서만 오차가 작고 나머지는 크면 Sinkhorn 반복이 흔들릴 수 있는데, 모든 행과 모든 열에 대해 상한이 걸리면 이야기가 달라집니다. 생략이 전체 수렴을 깨뜨리지 않는다는 쪽으로 논증할 수 있기 때문입니다.

이 보장이 중요한 이유는 생략이 매 반복마다 일어나기 때문입니다. 한 번만 건너뛰는 것이 아니라 정밀 단계의 매 반복에서 같은 스크리닝이 돌아가는데, 이때 오차가 쌓이면 주변 분포가 다시 어긋나기 때문입니다. 그래서 상한은 한 번의 편안함이 아니라 반복 전체의 안전장치입니다. 어느 블록을 버려도 행 합과 열 합의 어긋남이 어느 수준을 넘지 않는다는 약속이 있어야 반복을 끝까지 돌릴 수 있습니다.

남은 블록의 계산도 허투루 돌리지 않습니다. 논문은 fused tensor-core 실행으로 묶인 블록을 처리한다고 밝히는데, 이는 남긴 블록을 GPU에서 되도록 비싸지 않게 먹겠다는 선택입니다. 블록이 메모리에서 흩어져 있으면 가져오는 비용이 커지기 때문에, Morton 순서로 모은 배치는 여기서도 도움이 됩니다. 가까이 모인 점들의 블록을 통째로 올려서 tensor core로 밀어버리면, 남긴 계산의 단가를 낮출 수 있기 때문입니다. 저는 이 대목에서 이 논문이 알고리즘 논문이면서 동시에 시스템 논문이라고 느꼈습니다. 무엇을 계산하지 않을지와 남은 것을 어떻게 싸게 먹을지를 함께 설계했기 때문입니다.

32개 중 32개와 32개 중 10개는 어떻게 읽어야 할까

합성 벤치마크에서 나온 숫자는 이렇게 정리됩니다. 목표 잔차에 도달한 문제 수가 FlashSinkhorn 2는 32개 중 32개였고, 비교 대상인 GeomLoss multiscale은 32개 중 10개였습니다. 일반 파일에만 남겼을 때와 달리 이 숫자는 조건을 붙여서 읽어야 하는데, 같은 목표 잔차를 두고 누가 끝까지 도달했는지를 잰 것이기 때문입니다. 도달했는지 못했는지를 세는 방식이라서, 평균 오차를 비교하는 것과는 결이 다릅니다.

그렇다면 이 차이를 어떻게 해석해야 할까요. 저는 이 숫자보다 그 뒤의 구조가 더 중요하다고 봅니다. 성긴 단계만으로는 지워지지 않는 오차를 블록 희소 단계가 마저 지우는 구조라서, 어려운 문제에서도 잔차가 끝까지 내려간 것으로 읽기 때문입니다. 반대로 말하면, 성긴 단계에서 멈추거나 생략에 보장이 없는 설계는 몇몇 문제에서 잔차가 멈춰버릴 수 있다는 뜻이기도 합니다. 10개라는 숫자는 그 멈춤이 드물지 않았다는 신호로 볼 수 있습니다.

여기서는 조금 멈칫할 필요가 있습니다. 합성 문제는 설계하기에 따라 결과가 달라질 수 있기 때문입니다. 어떤 분포 모양을 넣었는지, 어려운 조건을 얼마나 빡빡하게 잡았는지에 따라 성공 개수는 움직입니다. 그래서 32대 10이라는 점수 자체를 일반 성능 우열로 단정하기보다는, 목표 잔차까지 밀어붙이는 마무리가 다르다는 쪽으로 읽는 것이 안전합니다. 특히 주변 분포 잔차를 끝까지 낮춰야 하는 용도에서는 마무리가 곧 쓸모이기 때문입니다.

1억 3천만 개 입자를 한 장의 GPU에 올린다는 것

대규모 실험은 우주론 N-body 시뮬레이션에서 나온 두 개의 입자 분포를 쓰는 것으로 되어 있습니다. 한쪽과 다른 쪽이 각각 1억 3천 4백만 개 수준의 입자를 갖고 있고, 이를 한 장의 A100에서 2시간 30분 안에 풀었다는 것이 논문의 주장입니다. 이 규모에서는 행렬을 만들 생각 자체가 불가능해서, 성긴 단계와 블록 희소 단계를 빼면 매 반복의 연산량을 감당하기 어렵습니다. 그래서 이 실험은 단순한 과시가 아니라 설계가 버티는지를 보는 시험에 가깝습니다.

조건도 함께 봐야 합니다. 엔트로피 흐림의 크기를 입자 사이 평균 거리와 같게 잡았고, 모든 입자에 대한 주변 잔차를 0.01 아래로 낮췄다는 것이 함께 보고되어 있습니다. 흐림이 크면 문제는 쉬워지지만 해가 뭉개지고, 흐림이 작으면 해는 날카로워지지만 계산이 어려워지기 때문입니다. 평균 거리와 같은 수준으로 잡았다는 것은 해상도와 계산 난이도 사이에서 실제 분포의 스케일에 맞춘 선택으로 읽힙니다. 그리고 잔차 0.01은 일부 표본이 아니라 모든 입자에 대해 잰 것이라는 점에서 마무리의 기준이 분명합니다.

좋은데, 이걸 어디에 올릴까요. 저는 이 숫자를 보고 메모리보다 시간과 대역폭을 먼저 떠올렸습니다. 한 장의 GPU라는 조건은 메모리가 넉넉하다는 뜻이 아니라, 오히려 메모리와 연산과 데이터 이동을 모두 한 장 안에서 감당했다는 뜻이기 때문입니다. 셀 중심점 단계가 전체 윤곽을 싸게 잡고, Morton 순서로 모인 블록만 tensor core로 처리하면서, 생략된 부분의 오차는 임계값으로 묶어두는 구조라야 이 시간이 나옵니다. 어느 하나가 빠져도 2시간 30분은 깨지기 쉽습니다. 밤사이 돌려야 하는 파이프라인이라면 이 차이는 그대로 마감 차이로 이어집니다.

그래서 다음에 제가 보고 싶은 장면은

이 논문이 남긴 질문을 정리하면 이렇게 됩니다. 성긴 단계가 잡은 윤곽 위에 정밀 단계가 마무리를 얹는 구조는 저차원 점구름에서는 잘 돌았는데, 차원이 올라가면 같은 스크리닝이 얼마나 버틸까요. 차원이 높아지면 거리만으로 near와 far를 가르기 어려워지고, Morton 순서의 군집도 흐릿해지기 때문입니다. 제곱 유클리드 비용과 저차원이라는 조건이 빠지면 블록 희소의 이득이 얼마나 남을지가 궁금합니다. 논문이 겨냥한 범위가 분명한 만큼, 그 밖에서는 다른 설계가 필요할 수도 있습니다.

또 하나 보고 싶은 것은 표본 검사의 거동입니다. 성긴 단계에서 들어 올리기를 받아들일지 말지를 표본으로 정하는데, 분포가 치우쳐 있거나 셀 크기가 고르지 않으면 표본이 실제 오차를 과소평가할 수 있기 때문입니다. 통과했는데 정밀 단계에서 오래 걸리는 경우가 얼마나 되는지, 중심점으로 돌아가는 횟수가 문제에 따라 어떻게 달라지는지가 수치로 나오면 설계의 감각을 더 잡을 수 있습니다. 운영 측면에서 보면 이 횟수는 그대로 대기 시간의 분산으로 이어지기 때문입니다.

마지막으로 남는 생각은 평가의 방향에 대한 것입니다. 목표 잔차에 도달한 문제 수를 세는 방식은 마무리를 보는 데는 좋지만, 같은 시간 안에 어디까지 내려가는지를 보는 것과는 답이 다를 수 있습니다. 제한된 예산에서 잔차를 얼마나 빨리 낮추는지, 흐림 크기를 바꾸면 시간과 정확도가 어떻게 맞바뀌는지가 함께 나오면 선택이 쉬워집니다. 제가 앞으로 보고 싶은 것은 바로 그 교환 곡선입니다. 1억 규모를 풀었다는 사실보다, 그 규모에서 무엇을 포기하고 무엇을 지켰는지가 다음 설계를 정하기 때문입니다.