🤖 machine learning

A Faster Generalized Two-Stage Approximate Top-K

본 논문은 각 파티션에서 상위 1 개가 아닌 상위 $K'$ 개 요소를 선택함으로써 2 단계 근사 Top-K 알고리즘을 일반화하여 더 엄밀한 이론적 재현율 상한을 제공하고 동일한 기대 재현율을 유지하면서 Cloud TPUv5e 에서 10 배의 속도 향상을 입증한다.

원저자: Yashas Samaga, Varun Yerram, Spandana Raj Babbula, Prateek Jain, Praneeth Netrapalli

게시일 2026-05-14

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Yashas Samaga, Varun Yerram, Spandana Raj Babbula, Prateek Jain, Praneeth Netrapalli

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

거대한 도서관의 관리자가 되어 상상해 보세요. 수백만 권의 책 (데이터) 이 있습니다. 매일 방문객들에게 추천할 가장 인기 있는 K 권의 책 (가장 큰 K 개의 숫자) 을 찾아야 합니다.

거대 AI 모델을 훈련시키는 데 사용되는 컴퓨터 칩 (특히) 의 세계에서는 이러한 '가장 인기 있는' 항목을 찾는 것이 놀랍도록 느리고 비용이 많이 듭니다. 도서관이 거대한 책 더미 전체를 한 번에 수학적으로 처리하도록 설계되어 있음에도 불구하고, 모든 책을 하나씩 읽어가며 상위 100 권을 찾으려 하는 것과 같습니다.

이 논문이 그 문제를 해결하기 위해 무엇을 하는지 간단히 설명해 드리겠습니다.

구식 방법: '한 번에 하나씩' 필터링

이전 방법 (Chern 등, 2022) 은 두 단계 프로세스를 사용하여 이를 가속화하려 했습니다:

분할: 도서관을 100 개의 다른 방 (버킷) 으로 나눈다고 상상해 보세요.
첫 번째 스캔: 각 방에서 조수가 단 한 권의 가장 인기 있는 책만 골라 프런트 데스크로 가져옵니다.
최종 정렬: 관리자는 이제 각 방에서 나온 100 권의 책 (각 방에서 하나씩) 만 보고 전체 상위 100 권을 선택합니다.

문제점: 이 방법은 너무 신중했습니다. 각 방에서 단 한 권의 최우수 책만 선택함으로써, 같은 방에 숨어 있던 두 번째나 세 번째로 좋은 책들을 종종 놓쳤습니다. 아무것도 놓치지 않도록 하기 위해 많은 방 (버킷) 을 사용해야 했으며, 이는 결국 관리자가 여전히 방대한 책 더미를 정렬해야 했다는 것을 의미했습니다. 여전히 너무 느렸습니다.

새로운 아이디어: 'Top-K' 필터

이 논문의 저자들은 컴퓨터 칩이 사용하지 않는 추가적인 성능이 있음을 깨달았습니다. 그들은 첫 번째 단계의 더 지능적인 버전을 제안했습니다:

각 방에서 #1 책만 고르는 대신, 조수가 이제 각 방에서 Top-K' 권의 책 (예를 들어 상위 4 권) 을 선택합니다.

왜 이것이 더 나은가요?

더 적은 방 필요: 조수가 각 방에서 더 많은 책을 가져오기 때문에, 모든 인기 있는 책을 확보하기 위해 필요한 방의 수가 줄어듭니다.
덜 많은 정렬: 조수가 방당 더 많은 책을 가져오지만, 최종 정렬을 위해 관리자에게 보내지는 책의 총 수는 실제로 훨씬 적습니다.
결과: 관리자는 산더미 같은 책 대신 아주 작은 더미를 정렬하면 됩니다.

하드웨어의 '마법'

이 논문은 현대 컴퓨터 칩 (예: Google 의 TPU) 이 다양한 작업 스테이션을 가진 거대한 공장이라고 설명합니다:

행렬 단위 (MXU): 무거운 수학 (곱셈) 을 수행하지만 정렬에는 서툰 초고속 공장.
벡터 단위 (VPU): 정렬과 승자 선출에 능숙하지만 더 작고 느린 작업 스테이션.

구식 방법은 VPU 의 시간을 낭비했습니다. 새로운 방법은 MXU 가 수학 계산을 하는 동안 VPU 가 'Top-K' 권의 책을 가져오도록 합니다. 이는 기계가 여전히 작동하는 동안 작업자가 컨베이어 벨트에서 최고의 물건을 가져오는 것과 같아 대기 시간이 없습니다.

결과: AI 가속화

저자들은 Google TPU 칩에서 이를 테스트했습니다:

구식 방법: 상위 책을 찾는 데 오랜 시간이 걸렸으며, 종종 처음에 목록을 생성한 수학 계산보다 더 느렸습니다.
신식 방법: 각 버킷에서 'Top 1' 대신 'Top 4'를 가져옴으로써 최종 정렬에 필요한 작업을 평균 7 배 줄였습니다.
퓨전: 그들은 '선택' 단계와 '수학' 단계를 결합하여 정확히 같은 시간에 발생하도록 만들었습니다.

핵심 결론:
실제 테스트 (대규모 AI 모델에서 상위 2% 의 데이터를 찾는 것) 에서 그들의 새로운 방법은 이전 표준보다 24 배 빠른 속도를 기록했습니다. 이는 AI 모델이 정확도를 잃지 않고 훨씬 더 빠르게 훈련하고 실행할 수 있음을 의미합니다.

요약 비유

구식 방법: 1,000 개의 팀이 있습니다. 각 팀은 당신에게 그들의 최고 선수를 보냅니다. 그런 다음 상위 100 명을 찾기 위해 1,000 명의 선수를 인터뷰해야 합니다.
신식 방법: 팀 수가 더 적습니다 (예: 250 개). 각 팀은 당신에게 그들의 상위 4 명의 선수를 보냅니다. 당신은 1,000 명의 선수 (250 개 팀 × 4 명) 만 인터뷰하면 되지만, 각 팀에서 더 많은 옵션을 받았기 때문에 진정한 최고의 선수를 찾을 확률은 동일하며 팀을 더 잘 조직했기 때문에 훨씬 빠르게 수행합니다.

이 논문은 수학적으로 이 'Top-K'' 접근 방식이 단순한 추측이 아니라, 훨씬 적은 작업으로 동일한 품질의 결과를 보장하는 방법임을 증명합니다.

기술 요약: 더 빠른 일반화된 2 단계 근사 Top-K

문제 정의
배열에서 상위 $K$ 개 요소를 식별하는 작업은 최대 내적 검색 (MIPS), K-최근접 이웃 (KNN), 그리고 대규모 기반 모델의 최적화 (예: 분류 로짓의 희소성, MLP 블록, 어텐션 메커니즘, KVCache 압축을 위한 희소성) 를 포함한 많은 머신러닝 알고리즘의 기본 연산입니다. TPU 및 GPU 와 같은 현대 가속기에서 정확한 Top-K 를 계산하는 것은 종종 성능 병목 현상이 됩니다. 이러한 장치는 밀집 행렬 곱셈 (MXU/TensorCores) 에 최적화되어 있는 반면, Top-K 연산은 일반적으로 스칼라/벡터 유닛 (VPUs/CUDA 코어) 에 의존하며 데이터를 생성하는 행렬 곱셈보다 훨씬 느릴 수 있습니다. 예를 들어, TPUv5e 에서 Gemma 2 9B 모델의 숨겨진 활성화 중 상위 2% 를 찾는 데는 해당 행렬 곱셈보다 27 배 더 많은 시간이 소요됩니다.

근사 Top-K 알고리즘은 대안을 제공하지만, Chern 등 (2022) 의 최첨단 방법인 jax.lax.approx_max_k 로 구현된 방법은 여전히 상당한 지연 시간을 초래합니다. 그들의 2 단계 접근법은 입력을 버킷으로 분할하고 각 버킷에서 상위 1 개 요소를 선택한 후 결과를 정렬합니다. 그러나 두 번째 단계 (선택된 요소 정렬) 는 여전히 병목 현상이며, 첫 번째 단계는 종종 사용 가능한 스칼라/벡터 연산 자원을 충분히 활용하지 못합니다.

방법론
저자들은 Chern 등 (2022) 의 알고리즘을 일반화하여 첫 번째 단계에서 상위 1 개로 제한하는 대신 각 버킷에서 상위 $K'$ 개 요소를 선택하도록 수정한 알고리즘을 제안합니다 ( $1 \le K' \le K$ ).

알고리즘 일반화:
- 1 단계: 크기 $N$ 인 입력 배열을 $B$ 개의 버킷으로 분할합니다. 최대값만 선택하는 대신, 알고리즘은 각 버킷에서 상위 $K'$ 개 요소를 선택합니다.
- 2 단계: 선택된 $B \times K'$ 개의 요소를 병합하여 최종 상위 $K$ 개 요소를 반환하기 위해 정렬합니다.
- 구현: 첫 번째 단계는 TPU 를 위한 JAX 커널 언어인 Pallas 를 사용하여 온라인 업데이트 서브루틴으로 구현됩니다. 각 버킷에 대해 상위 $K'$ 개 값과 인덱스의 정렬된 목록을 유지합니다. 새로운 요소는 버킷 간 벡터화를 최적화한 버블 정렬과 유사한 패스를 사용하여 삽입됩니다. 저자들은 두 개의 별도 커널을 사용하는 "언퓨즈드 (unfused)" 버전과 첫 번째 단계를 이전 행렬 곱셈과 병합하여 첫 번째 단계 비용을 효과적으로 숨기는 "행렬 곱셈 퓨즈드 (matmul-fused)" 버전을 모두 제공합니다.
이론적 분석:
- 저자들은 무작위 분할 하에서 이 일반화된 알고리즘의 기대 재현율에 대한 정확한 식을 유도하며, 하이퍼기하 분포를 사용하여 상위 $K$ 개 요소를 버킷에 분포시키는 모델을 제시합니다.
- 고정된 기대 재현율 목표에 대해 $K'$ 을 증가시키면 두 번째 단계로 전달되는 요소의 총 수 ( $B \times K'$ ) 를 최소화할 수 있도록 버킷 수 $B$ 를 줄일 수 있음을 증명합니다.
- 특히, $K'=1$ 일 때 필요한 버킷 수에 대한 새로운 상한선을 유도합니다. 이 상한선은 Chern 등 (2022) 의 상한선보다 2 배 더 엄격하게 증명되어, 원래 설정에서도 더 적은 버킷과 따라서 정렬해야 할 요소가 적어지도록 합니다.
성능 모델링:
- 릿지 포인트 분석을 사용하여 저자들은 가속기의 하위 시스템 (HBM, VPU, MXU) 을 특성화합니다. 그들은 첫 번째 단계가 일반적으로 $K'$ 이 특정 임계값 (예: TPUv5e 에서 $K' \approx 6$ ) 에 도달할 때까지 메모리 병목 현상을 겪음을 보여줍니다. 이 메모리 병목 구간 내에서 $K'$ 을 증가시켜도 첫 번째 단계의 지연 시간은 크게 증가하지 않지만 두 번째 단계의 작업량은 극적으로 감소합니다.

주요 기여

일반화된 알고리즘: 각 버킷당 상위 $K'$ 개 요소를 선택하는 새로운 2 단계 근사 Top-K 알고리즘으로, 첫 번째 단계 연산과 두 번째 단계 정렬 부하 사이의 유연한 트레이드오프를 제공합니다.
이론적 상한선:
- $N, K, B,$ 및 $K'$ 의 함수로서 기대 재현율에 대한 정확한 식.
- 이전 작업에 비해 $K'=1$ 인 경우 필요한 버킷 수에 대해 2 배 더 엄격하게 증명된 상한선으로, 매개변수 선택을 개선합니다.
구현: 퓨즈드 및 행렬 곱셈 퓨즈드 변형을 특징으로 하는 Cloud TPUv5e 에서의 고성능 Pallas 구현.
자동 매개변수 선택: 사용자가 지정된 재현율 목표를 충족하도록 $K'$ 과 $B$ 를 자동으로 선택하는 루틴으로, 수동 튜닝의 필요성을 제거합니다.

결과

이론적 효율성: 고정된 재현율 목표 (예: 95%) 에 대해 $K'$ 을 증가시키면 두 번째 단계에서 처리되는 요소 수가 크게 감소합니다. 예를 들어, 262,144 개 요소 중 상위 1024 개를 선택하는 경우 $K'=1$ 에서는 16,384 개의 출력 요소가 필요하지만 $K'=4$ 에서는 2,048 개 요소만 필요합니다 (8 배 감소).
지연 시간 개선 (언퓨즈드): TPUv5e 에서 $K'=4$ 인 일반화된 알고리즘은 동일한 작업에서 99% 재현율 목표에 대해 원래 jax.lax.approx_max_k 대비 11 배의 속도 향상 (27µs 대 305µs) 을 달성합니다.
지연 시간 개선 (퓨즈드): 최대 내적 검색 (MIPS) 워크로드 (100 만 개 벡터, 1024 개 쿼리) 에서 $K'=4$ 인 행렬 곱셈 퓨즈드 구현은 총 실행 시간을 594ms(정확한 Top-K) 에서 10ms 로 줄입니다. 이는 밀집 기준 대비 6.7 배의 전체 속도 향상으로, 근사 Top-K 단계를 행렬 곱셈 자체의 비용보다 낮게 만듭니다.
종단 간 영향: 희소 트랜스포머 블록 벤치마크 (Gemma 2 9B) 에서 제안된 방법은 원래 알고리즘을 사용할 때 밀집 대비 2.7 배 느렸던 희소 MLP 블록의 오버헤드를 밀집 기준 대비 불과 +5ms 의 소폭 오버헤드로 줄입니다.

의의 및 주장
이 논문은 이전 근사 Top-K 알고리즘에서 $K'=1$ 을 선택하는 것이 현대 가속기의 사용 가능한 연산 자원을 완전히 활용하지 못한다고 주장합니다. $K' > 1$ 로 일반화함으로써, 알고리즘은 첫 번째 단계가 메모리 병목 현상에 머무는 한 필터링 단계에서 상당한 추가 비용을 치르지 않고 병목 현상을 정렬 단계에서 효과적으로 이동시킵니다.

저자들은 이론적 분석이 매개변수 선택을 위한 엄격한 기초를 제공하며, 수동 튜닝이 필요 없는 사용자 친화적인 인터페이스 (approx_top_k(array, K, recall_target)) 를 제공한다고 강조합니다. 실증적 평가는 TPUv5e 에 초점을 맞추고 있지만, 저자들은 알고리즘, 성능 모델링 및 분석이 일반적이며 GPU 와 같은 다른 하드웨어 플랫폼에 적용 가능하다고 주장합니다. 이 작업은 hardware 인식 구현 (퓨전) 과 결합된 원칙적인 알고리즘 일반화가 대규모 모델 학습 및 추론에서 Top-K 를 병목 현상으로 제거할 수 있음을 보여줍니다.