🤖 AI

Fast LapSum: Exact Differentiable Top-k at Million Scale

이 논문은 $k$ 라는 정확한 선택 질량을 보존하면서 GPU에서 선형 시간 내에 실행되는 정밀하고 미분 가능한 소프트 톱- $k$ 프리미티브인 Fast LapSum을 소개하며, 이를 통해 적대적 예제 생성 및 미분 가능한 이미지 코딩과 같은 응용 분야를 위한 백만 단위 규모의 희소 연산을 효율적으로 가능하게 한다.

원저자: Łukasz Struski, Joanna Wojciechowicz, Jakub Antczak, Marcin Mazur, Kamil Książek, Jacek Tabor

게시일 2026-08-10

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Łukasz Struski, Joanna Wojciechowicz, Jakub Antczak, Marcin Mazur, Kamil Książek, Jacek Tabor

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

당신이 매초 수백만 권의 책이 스캔되고 있는 거대한 디지털 도서관을 운영하고 있다고 상상해 보십시오. 이 정보의 홍수를 이해하기 위해, 도서관의 AI는 지금 당장 가장 중요한 책 몇 권이 무엇인지 결정해야 합니다. 인공지능의 세계에서 이것은 "top-k 선택(top-k selection)"이라고 불립니다. 즉, 거대한 목록에서 가장 좋은 $k$ 개의 항목을 골라내는 것입니다. 보통 AI는 상위권의 책들만 골라내고 나머지는 완전히 무시하는 엄격한 사서처럼 행동합니다. 이는 속도 면에서는 훌륭하지만, 학습 측면에서는 최악입니다. 왜냐하면 AI가 어떻게 더 나아질 수 있는지 알아낼 방법이 없기 때문입니다. 이는 마치 이미 올바른 차선에 있을 때만 도로를 보고 운전하는 법을 배우려는 것과 같아서, 핸들을 조절하는 법을 배울 방법이 없는 것과 같습니다.

이를 해결하기 위해 과학자들은 "소프트(soft)" 버전의 선택법을 발명했습니다. "예 또는 아니오"라는 단호한 결정 대신, AI에게 모든 책에 대해 "아마도"라는 점수를 부여하여 실수를 통해 배울 수 있게 하는 것입니다. 하지만 여기에는 함정이 있습니다. 이러한 소프트 버전들은 종 Часто 너무 느리고 계산량이 많아서 도서관이 커지면 시스템을 다운시켜 버립니다. 이는 마치 도서관에 불이 났는데 수백만 권의 책을 손으로 직접 분류하려는 것과 같습니다. 연구자들의 큰 질문은 이것이었습니다. "학습할 수 있을 만큼 부드러우면서도(미분 가능하면서도), 수백만 권의 책을 처리할 수 있을 만큼 빠른 라이브러리를 만들 수 있을까?"

여기서 새로운 논문인 "Fast LapSum"이 등장합니다. 저자들인 폴란드 팀은 수학적으로 완벽하고 매우 효율적인 사서 역할을 하는 새로운 도구를 구축했습니다. 그들은 Fast LapSum이라는 방법을 만들어, AI가 수백만 개의 목록 중에서 상위 항목들을 골라내면서도 그 과정을 통해 학습할 수 있도록 했습니다. 기존의 방식들이 속도를 얻기 위해 완벽한 정확성을 포기하거나, 혹은 너무 느려서 쓸모가 없었던 것과 달리, Fast LapలSum은 이 두 가지를 모두 해냅니다. 이 방법은 뽑아야 할 정확한 아이템의 개수(예산)를 찾아내고, 그들에 대한 완벽한 "아마도" 점수를 눈 깜짝할 사이에 계산합니다.

그 비결은 점수의 "흐릿한(blurred)" 뷰를 이용한 영리한 트릭에 있습니다. 점수가 날카로운 점이 아니라 흐릿한 구름이라고 상상해 보십시오. AI는 이 구름들의 총량이 자신이 허용된 책의 수와 정확히 일치하도록 선을 그려야 합니다. 기존 방식들은 이 선을 찾기 위해 계속해서 추측하고 확인하는 과정을 반복했기에 시간이 너무 오래 걸렸습니다. 그러나 Fast LapSum은 "라플라스 분포(Laplace distribution)"라고 불리는 특별한 수학적 공식을 사용하여, 단 한 번의 정렬(sort)만으로도 즉시 선을 계산할 수 있게 해줍니다.

백만 개 또는 천만 개의 점수와 같이 정말 거대한 목록의 경우, 저자들은 "확률적 브래키팅(probabilistic bracketinging)"이라는 두 번째 트릭을 추가했습니다. 경기장 가득 찬 사람들을 정리하는 대신 전체 목록을 정렬하는 대신, 시스템은 선이 아마도 어디쯤 있을지 추측하기 위해 빠른 샘플링을 수행합니다. 그런 다음 그 선 근처에 서 있는 작은 그룹만을 정렬합니다. 이를 통해 프로세스를 믿을 수 없을 정도로 빠르게 유지하며, 거대한 데이터셋에 대해서도 단 몇 밀리초 만에 처리가 가능합니다.

논문은 두 가지 어려운 과제를 통해 이 방법이 작동함을 증명합니다. 첫째, 인간에게는 정상적으로 보이지만 AI 분류기를 속이는 이미지인 "적대적 예시(adversarial examples)"를 만드는 데 사용되었습니다. 그들은 이미지를 아주 미세하게 변형하여(330만 픽셀 중 약 600 픽셀, 즉 0.02% 정도만 변경) AI가 호랑이 사진을 오인하게 만들었습니다. 이는 이전 방식들보다 훨씬 빠르고 이미지에 주는 "손상"도 적었습니다. 둘째, 가장 중요한 부분만을 남기도록 이미지를 압축하는 시스템인 미분 가능한 이미지 코더(differentiable image coder)를 처음부터 구축했습니다. 두 경우 모두 Fast LapSum은 엔진 역할을 하며, 학습 과정을 늦추지 않고 초당 수백만 개의 결정을 처리했습니다.

저자들은 이 방법이 단순한 이론적 아이디어가 아니라 표준 컴퓨터 칩에서 밀리초 단위로 실행되는 실용적인 도구임을 보여줍니다. 그들은 DFTopK와 같은 최근의 다른 시도들과 비교하였으며, 그 방법들이 빠르기는 하지만 선택의 정확성(뽑힌 아이템의 총량이 목표치에서 벗어남)을 희생한다는 것을 발견했습니다. Fast LapSum은 선택을 완벽하게 정확하게 유지하면서도 실제 대규모 AI 시스템에 적용될 수 있을 만큼 빠르다는 점에서 최초의 방식이라고 주장합니다. 이 방법은 느리고 비용이 많이 드는 병목 현상을 매끄럽고 빠른 작업으로 바꾸어 놓음으로써, AI가 똑똑하면서도 효율적일 수 있게 만듭니다.

기술 요약: Fast LapSum: 백만 규모에서의 정밀한 미분 가능한 Top-k

1. 문제 정의

top-k 연산은 현대의 희소 연산(sparse computation)에서 필수적인 기초 프리미티브로, Mixture-of-Experts (MoE) 모델의 토큰 라우팅, 롱 컨텍스트 아키텍처의 메모리 선택, 그리고 검색 시스템 등에서 핵심적인 역할을 합니다. 그러나 표준적인 하드 top-k 선택 방식은 미분이 불가능하여 결정 경계(decision boundary)를 가로지르는 그래디언트 흐름을 차단합니다. 반대로, 기존의 연속적(soft) 완화 방식(예: NeuralSort, SoftSort, Sinkhorn-OT)은 과도한 계산 비용( $O(n^2)$ 또는 반복적 복잡도)을 초래하거나, 정확한 선택 예산(선택된 확률의 합이 $k$ 가 되어야 함)을 유지하지 못하는 경우가 많습니다.

현재의 미분 가능한 top-k 방법들은 트레이드오프에 직면해 있습니다: 이들은 대규모 모델의 학습 루프 내에서 실행하기에는 너무 비싸거나, 정규화 제약 조건( $\sum p_i = k$ )을 완화하여 통제되지 않는 희소성을 초래합니다. 본 논문은 미분 가능한 희소화 도구(sparsifier)가 하드 top-k의 정확한 예산을 유지하면서도 백만 규모의 벡터에서도 계산적으로 충분히 저렴한 "시스템 프리미티브"가 되어야 한다고 주장합니다.

2. 방법론

2.1 LapSum 연산자

저자들은 top-k 선택을 확률적으로 모델링하는 LapSum 연산자를 기반으로 구축하였습니다. 점수 $r_i$ 를 결정론적인 이진 마스크 대신, 라플라스 노이즈 $\epsilon \sim \text{Laplace}(0, t)$ 에 의해 섭동(perturbed)된 것으로 취급합니다. 선택 확률 $p_i$ 는 섭동된 점수가 스칼라 임계값 $b$ 를 초과할 확률로 정의됩니다:
$p_i = \sigma\left(\frac{r_i - b}{t}\right)$
여기서 $\sigma$ 는 라플라스 누적 분포 함수(CDF)입니다. 임계값 $b$ 는 총 질량이 예산 $k$ 와 같아야 한다는 제약 조건에 의해 유일하게 결정됩니다:
$\sum_{i=1}^n p_i = k$
이전 방식들이 이 임계값을 근사하는 것과 달리, LapSum은 $b$ 를 정확하게 계산합니다. 역전파(backward pass)는 편향된 대리물인 straight-through estimator를 피하는 벡터-자코비안 곱(VJP)을 활용하여 해석적으로 수행됩니다.

2.2 Fast LapSum: GPU 솔버

핵리의 기여는 정렬 후 선형 시간 내에 정확한 임계값 $b$ 를 계산하는 GPU 네이티브 솔버인 Fast LapSum입니다.

Full-Sort 솔버: 정렬이 가능한 벡터의 경우, 이 방법은 정렬된 점수들에 대해 단일 안정적 prefix/suffix 스캔을 사용하여 "임계값 곡선(threshold curve)"을 계산합니다. 임계값 $b$ 는 인덱스 룩업과 예산 방정식이 균형을 이루는 특정 구간 내에서의 폐쇄형 루트(closed-form root) 계산을 통해 찾아집니다. 이는 이전 구현에서 사용되던 이분법(bisection)을 대체하여, 정렬 이후의 작업을 $O(n)$ 으로 줄여줍니다.
확률적 브래키팅 (백만 규모): 전체 벡터를 정렬하는 것이 병목이 되는 극단적인 규모(예: $n > 10^6$ $n > 1 0^{6}$ )를 위해, 저자들은 확률적 브래키팅 파이프라인을 도입했습니다:
1. 국소화(Localization): 샘플링된 커널 노이즈 점수를 사용하여 예산에 해당하는 분위수(quantile)를 추정합니다. 이항 순서 통계량(binomial order statistics)을 사용하여 실제 임계값 $b$ 를 높은 확률로 포함하는 좁은 구간 $[a_L, a_R)$ 을 식별합니다.
2. 스트리밍 및 필터링: 전체 벡터를 한 번 스트리밍합니다. 구간 외의 점수들은 라플라스 CDF 특성에 따라 안정적인 꼬리 합(stable tail sums, $\tilde{S}_L, \tilde{S}_R$ )으로 집계됩니다. $[a_L, a_R)$ 범위 내의 "불확실한 중간 밴드(uncertain middle band)" 점수들만 유지됩니다.
3. 정확한 해 도출: 정확한 LapSum 솔버를 이 축소된 중간 밴드에 대해서만 적용합니다.
4. 검증: 브래킷이 유효한지 확인하는 인증 단계가 있으며, 유효하지 않을 경우 반복적으로 범위를 넓힙니다. 이를 통해 최종 결과가 정확한 예산 제약 조건 $\sum p_i = k$ 를 만족함을 보장합니다.

2.3 DFTopK와의 관계

본 논문은 Fast LapSum을 동시기에 발표된 DFTopK와 구분합니다. DFTopK는 임계값을 $k$ 번째와 $(k+1)$ 번째 순서 통계량의 중간값으로 고정함으로써 선형 시간을 달성하지만, 이는 정규화 제약을 완화하는 방식입니다. 저자들은 DFTopK가 LapSum의 제로 템퍼러처(zero-temperature) 극한임을 증명합니다. 그러나 유한한 온도에서 Fast LapSum은 정확한 예산을 보존하는 반면, DFTopK는 활성 질량이 드리프트(예: $k$ 대신 $\approx 3.46k$ 의 질량 발생)하는 것을 허용합니다.

3. 주요 기여

Fast LapSum 연산자: 정확한 예산 제약 $\sum p_i = k$ 를 엄격히 준수하면서 선형 정렬 후 복잡도를 달성한 최초의 GPU 솔버입니다. 이는 편향된 대리물을 제거합니다.
확률적 브래키팅 파이프라인: 샘플링을 통해 임계값을 국소화하고, 불확실한 중간 밴드만을 정렬하며, 정확한 해를 검증하여 기대 선형 시간 복잡도를 유지하는 백만 규모 벡터용 방법론입니다.
확장성 및 효율성: $10^8$ 개의 점수를 수 밀리초(예: $10^8$ 개 점수에 대해 5.23 ms) 내에 처리할 수 있는 능력을 입증하였으며, 이는 기존의 정확한 구현들보다 수 차례 더 빠릅니다.

4. 실험 결과

4.1 성능 벤치마크

속도: Fast LapSum의 full-sort 브랜치는 per-query 이분법 베이스라인보다 거의 두 자릿수 더 빠릅니다. 브래킷 브랜치는 $10^8$ 개의 점수에서 최대 19 $\times$ 의 속도 향상을, $3.2 \times 10^7$ 개의 점수에서 공식 LapSum 릴리스 대비 1200 $\times$ 의 속도 향상을 달상했습니다.
메모리: 밀집 정렬 완화 방식(NeuralSort, SoftSoft)이 $n=10^5$ 부근에서 메모리 부족(OOM)을 일으키는 것과 달리, Fast LapSum은 $n=10^8$ 에서도 효율적으로 작동합니다.
정확도: Fast Lapsum은 모든 규모에서 정확한 예산 정규화(상대 오차 $< 10^{-5}$ )를 유지합니다. 반면, DFTopK는 완화된 제약 조건으로 인해 약 2.46의 상대 오차(질량 $\approx 3.46k$ )를 보입니다.

4.2 응용 분야

저자들은 두 가지 까다로운 백만 규모 응용 분야에서 이 연산자를 시연했습니다:

미분 가능한 메가픽셀 이미지 코더: 처음부터 훈련된 완전 미분 가능한 희소 이미지 코더입니다. 이 연산자는 레이트-왜곡(rate-distortion) 루프 내에서 수백만 개의 점수에 대한 좌표 선택을 관리하며, 계산 비용을 지배하지 않으면서 희소 원자(atoms)로부터 이미지를 성공적으로 조립합니다.
희소 적대적 공격(Sparse Adversarial Attacks): ConvNeXt-V2 분류기를 대상으로 네이티브 해상도 이미지(DIV2K)에 대한 희소 적대적 예제를 생성하는 데 사용되었습니다.
- 효율성: 이 공격은 199.8의 중앙값 수정 픽셀 상당량의 예산으로 100% 성공률을 달성했습니다.
- 비교: 이는 베이스라인들보다 현저히 효율적입니다: $\sigma$ -zero는 3,101, Sparse-PGD는 16,134, PGD-L0는 74,553이 필요합니다.
- 정밀도: 이 공격은 흩어진 노이즈가 아니라 클래스를 정의하는 텍스처(예: 호랑이의 눈)에 섭동을 집중시키며, 이미지 픽셀의 약 $0.02\%$ 만을 사용합니다.

5. 의의 및 주장

본 논문은 Fast LapSum이 미분 가능성, 정확한 예산 제약, 그리고 계산 효율성 사이의 오랜 트레이드오프를 해결한다고 주장합니다. 미분 가능한 top-k를 시스템 프리미티브로 취급함으로써, 이 방법은 다음을 가능하게 합니다:

엔드 투 엔드 훈련: 이 연산자는 빠르기 때문에 대규모 희소 모델(라우팅, 검색, 프루닝)의 학습 루프 내에 직접 삽입되어도 계산 병목이 되지 않습니다.
정확한 희소성 제어: 질량이 드리프트되는 soft 완화 방식과 달리, Fast LapSum은 활성 전문가, 메모리 항목 또는 검색 후보의 수가 지정된 예산 $k$ 를 엄격히 준수하도록 보장합니다.
규모에서의 실용성: 이 방법은 고해상도 이미지 처리 및 대규모 검색과 같이 수백만 개의 좌표를 다루는 실제 작업에서 정확한 soft top-k를 실용적으로 만듭니다.

저자들은 제한 사항을 언급했습니다: 폐쇄형 해법은 라플라스 커널에 특화되어 있으며, 선형 스케일링은 최악의 경우(pathological score distributions가 추가 검증 단계를 유발할 수 있음)보다는 기대되는 성능입니다. 그러나 경험적 결과는 이 연산자가 규모가 커져도 무시할 수 있는 비용을 가지며, 정확한 예산 제약을 가진 새로운 클래스의 미분 가능한 시스템을 가능하게 함을 시사합니다.