Top-k retrieval에서 필요한 최소 임베딩 차원

Source

TL;DR

top-k retrieval에 필요한 최소 임베딩 차원(MED)은 아이템 수 m과 무관하게 2k면 충분하다는 이론적 증명.

기존 연구가 주장한 “m에 대한 polynomial 성장”은 기하학적 한계가 아니라 최적화 실패의 artifact였음
MED = Θ(k): inner product / euclidean / cosine 모두 scoring function에 무관하게 동일하게 성립
실제 병목은 learnability. 공간은 충분히 넓고, 문제는 올바른 벡터를 학습하는 것

Background

Embedding-based retrieval: universe $X$ 의 m개 원소를 $\mathbb{R}^d$ 에 임베딩, 쿼리 $q$ 도 벡터로 변환하여 scoring function $s(\cdot, \cdot)$ 으로 top-k 결과를 반환 한다고 할 때,
핵심 질문: m개 원소와 모든 top-k 쿼리를 완벽히 답하기 위한 최소 차원 d는 얼마인가?
이 최소 차원을 Minimal Embeddable Dimension (MED) 으로 정의
On the Theoretical Limitations of Embedding-Based Retrieval (2025) 에서는 시뮬레이션한 결과 최소 차원이 m에 대해 polynomial하게 성장한다고 주장
- $m^*(d) = -10.53 + 4.03d + 0.052d^2 + 0.0037d^3$ (Equation 1)
- -> 즉 m이 커질수록 필요한 d가 polynomial하게 증가 → 아이템이 많아질수록 차원을 엄청나게 늘려야 하고, 이는 벡터 공간의 근본적 한계라는 결론
이 논문은 그 결론을 이론적으로 반박 - MED는 m과 무관하게 2k차원이면 충분하고, 실제 병목은 learnability 문제임을 주장함 (기하 공간은 충분히 넓은데 그냥 딱 맞는 벡터를 뽑아내는 모델을 학습하는 게 어려운 거임)

$X \subseteq \mathbb{R}^d$ $X \subseteq R^{d}$ 가 k-shattered by $\mathcal{F}$ $F$ 이면:
- $\forall S \in \mathcal{C}_k, \exists f_S \in \mathcal{F}$ : $S$ 의 원소는 $b_S$ 보다 크고, $S$ 밖의 원소는 $b_S$ 보다 작은 score를 가짐
- 즉, 모든 top-k 쿼리를 완벽히 분리할 수 있는 임베딩 구성이 존재함을 의미

Definition (MED): m개 원소를 k-shattered 할 수 있는 최소 차원 $n^* = \text{MED}(m, k; \mathcal{F})$

$\text{VCD}^{-1}(m; \mathcal{F}) \leq \text{MED}(m, k; \mathcal{F}) \leq \text{VCD}^{-1}(m; \mathcal{F})$
$\text{MED}(m, m; \mathcal{F}) = \text{VCD}^{-1}(m; \mathcal{F})$ $MED (m, m; F) = VCD^{- 1} (m; F)$ (Lemma 2.7)
- -> VC dimension의 역함수가 MED와 연결됨

기존 연구에서는 아이템 벡터 m개와 쿼리 벡터 $\binom {m}{k}$ 개를 최적화했음 -> 쿼리 벡터가 문제(m이 늘수록 변수가 폭발적으로 많아지므로 optimizer가 좋은 해를 못 찾음)
쿼리 임베딩을 자유롭게 최적화하는 대신, answer set의 centroid로 고정: $w_q = \frac{1}{|S|}\sum_{x \in S} x$
자유도가 줄어 더 어려운 세팅이지만 시뮬레이션이 훨씬 용이 ( $\binom{m}{k}$ 개 쿼리 임베딩 최적화 불필요, m개만 최적화)
$\text{MED}(m, k; \mathcal{F}) \leq \text{MED-C}(m, k; s)$ (MED가 MED-C의 lower bound)

k-1 \leq \text{MED}(m, k; \mathcal{F}_{\text{linear}}) \leq 2k

Upper bound: Cyclic polytope 구성을 이용
- Moment curve $x(t) = (1, t, t^2, \ldots, t^d) \in \mathbb{R}^d$ 위의 점들로 이루어진 cyclic polytope
- $\mathbb{R}^d$ 의 cyclic polytope는 $\lfloor d/2 \rfloor$ -neighborly polytope: $k \leq \lfloor d/2 \rfloor$ 이면 모든 k개 꼭짓점이 face를 형성 → linear separability 보장
- 따라서 $d = 2k$ 이면 모든 top-k 쿼리를 분리 가능
Lower bound: $\mathcal{F}_{\text{linear}}$ 의 VC dimension이 $n+1$ 임을 이용 (Proposition 2.8)
핵심: MED는 m에 의존하지 않는다. 오직 k에 의존함

k-1 \leq \text{MED}(m, k; \mathcal{F}_{\ell_2}) \leq 2k

$\text{MED}(m, k; \mathcal{F}_{\ell_2}) \leq \text{MED}(m, k; \mathcal{F}_{\text{linear}})$ $MED (m, k; F_{ℓ_{2}}) \leq MED (m, k; F_{linear})$ (Proposition 3.3)
- Linear으로 k-shattered되는 구성이 있으면, 분리 hyperplane에 접하는 $\ell_2$ ball로 동일한 분리 가능

k-1 \leq \text{MED}(m, k; \mathcal{F}_{\cos}) \leq 2k+1

$\text{MED}(m, k; \mathcal{F}_{\text{linear}}) \leq \text{MED}(m, k; \mathcal{F}_{\cos}) \leq \text{MED}(m, k; \mathcal{F}_{\text{linear}}) + 1$ $MED (m, k; F_{linear}) \leq MED (m, k; F_{c o s}) \leq MED (m, k; F_{linear}) + 1$
- Cosine의 decision boundary = 구(sphere)와 hyperplane의 교선
- Linear → Cosine: inverse stereographic projection으로 $\mathbb{R}^n \to S^n \subset \mathbb{R}^{n+1}$
- 차원이 최대 1 증가

Probabilistic method 사용: $v_1, \ldots, v_m \sim \mathcal{N}(0, I_n/n)$ 을 랜덤 샘플링
두 벡터의 inner product 및 norm 집중 부등식: $\Pr\left[|\langle v_i, v_j \rangle| \geq \frac{1}{3k}\right] \leq 2\exp\left(-c\frac{n}{k^2}\right)$
Union bound 적용 시, $n > Ck^2 \log m$ $n > C k^{2} lo g m$ 이면 양의 확률로 k-centroid shattering 구성이 존재
- -> 즉 $\text{MED-C}(m, k; s) = O(k^2 \log m)$

역설적으로 자유도가 더 많은 free embedding optimization이 더 나쁜 결과를 냄
이유: free optimization에서 $\binom{m}{k}$ $(k m)$ 개의 쿼리 임베딩을 동시에 최적화해야 하는데, 이 landscape가 훨씬 어렵다 (local optima 등 최적화 문제)
- Centroid setting은 m개만 최적화 → 더 효과적으로 좋은 구성 탐색 가능
-> 즉 기존 연구의 polynomial 성장 주장은 최적화 실패의 artifact이지, 기하학적 근본 한계가 아님

결과: MED-C는 log-linear 성장 (empirically $\approx 3.23 \log_2 m$ $\approx 3.23 lo g_{2} m$ )
- 기존 연구의 cubic 성장 곡선을 쉽게 초과 → 훨씬 많은 원소를 같은 차원에 수용 가능