🤖 machine learning

Is Dimensionality a Barrier for Retrieval Models?

This paper resolves the theoretical question of why low-dimensional embeddings suffice for large-scale retrieval by proving that the optimal margin achievable in infinite dimensions can be nearly attained in dimension $d = O(k\log(n/k))$ for $k$ -sparse relevance matrices, while also empirically demonstrating the superiority of sigmoid loss over InfoNCE for generating large-margin embeddings.

Original authors: Kiril Bangachev, Guy Bresler, Jonathan Kogan, Yury Polyanskiy

Published 2026-05-25

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Kiril Bangachev, Guy Bresler, Jonathan Kogan, Yury Polyanskiy

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to organize a massive library containing billions of books. You want to find the right book for a specific question instantly. To do this, you create a "summary card" for every book and every possible question. These cards are just lists of numbers (vectors) that represent the content.

The big mystery this paper tackles is: How can these summary cards be so short and simple (low-dimensional) while still working perfectly for a library of trillions of items?

Usually, we think that to handle a huge, complex world, you need a huge, complex map. If you have billions of items, you'd expect the summary cards to need thousands or millions of numbers to be accurate. But in reality, modern AI systems use cards with only about 1,000 numbers and still find the right answers almost perfectly.

This paper asks: Is the small size of these cards a problem? Or is it actually a feature?

The Core Concept: The "Safety Margin"

The authors introduce a concept called the Margin. Think of this as a "safety buffer" or a "fence."

The Goal: You want to separate "relevant" books from "irrelevant" ones.
The Fence: Imagine drawing a line (or a wall) between the two groups.
The Margin: This is the distance from the books to that wall.
- If the margin is tiny, the books are right up against the wall. A tiny mistake (like a typo in the question or a smudge on the book) could push a book over the wall, and you'd pick the wrong one.
- If the margin is huge, there is a wide, safe zone. Even if the question is slightly different or the book is slightly different, it stays on the correct side of the wall.

The paper argues that a large margin is the secret to quality. It makes the system robust (it doesn't break easily) and generalizable (it can handle new, slightly different questions).

The Big Discovery: You Don't Need a Big Room

The authors wanted to know: How big does the room (the number of dimensions) need to be to build a fence with a huge safety margin?

The Old Belief: You probably need a massive room (high dimensions) to fit all the books and build a wide fence.

The Paper's Finding: You actually need a surprisingly small room.

They proved mathematically that you can achieve the best possible safety margin in a room that is only slightly larger than the logarithm of the number of books.
The Analogy: Imagine you have a billion books. You might think you need a stadium-sized room to organize them safely. The paper says, "Nope, a small, well-organized closet is enough." The size of the room only needs to grow slowly (logarithmically) as you add more books.

This explains why current AI models work so well with small vectors: The "low dimension" isn't a barrier; it's actually sufficient for the best possible performance.

The Two Main Experiments: The "Sigmoid" vs. "InfoNCE"

The researchers also tested two different ways of training these summary cards (two different "loss functions," which are like the rules the AI follows to learn).

InfoNCE: This is the popular method used by many current systems.
- The Result: It struggled. To get a positive safety margin (a working fence), it needed a much larger room (higher dimensions). It was like trying to build a fence in a crowded room; it kept bumping into things.
Sigmoid Loss: This is a different, slightly older method.
- The Result: It was a superstar. It built a perfect, wide safety margin in a tiny room. It succeeded where the other method failed, needing far fewer dimensions to get the job done.

The Takeaway: If you want your summary cards to be small and efficient, the "Sigmoid" method is the better architect.

Summary of the "Magic"

The Problem: Why do small, simple AI models work on huge datasets?
The Answer: Because you don't need a giant space to create a strong separation (margin) between good and bad answers.
The Proof: The authors used advanced math (connecting ideas from signal processing and geometry) to prove that the "best possible" separation can be achieved in a very small space.
The Practical Tip: If you are building these models, using the Sigmoid loss function helps you get that perfect separation in a much smaller, more efficient space than the standard methods.

In short: Small is beautiful. You don't need to make your data representations huge to get high-quality results; you just need the right mathematical tools to arrange them in a small space.

Technical Summary: Is Dimensionality a Barrier for Retrieval Models?

Problem Statement

Modern information retrieval systems rely on dense neural embeddings to represent data in low-dimensional spaces (typically $d \approx 1000$ or less) to enable efficient nearest-neighbor search. Despite the geometric limitations imposed by low dimensionality, these models achieve remarkable performance on massive datasets (billions to trillions of points). This paper investigates the theoretical underpinnings of this phenomenon: How can embedding-based retrieval models scale to huge datasets while maintaining high quality (robustness and generalization) in low dimensions?

The authors formalize this using a maximal-margin retrieval model. Given an incidence matrix $A \in \{0, 1\}^{N \times n}$ representing relevance between $N$ queries and $n$ documents, the goal is to find unit-norm embeddings $\{U_j\}$ and $\{V_i\}$ in dimension $d$ that maximize the margin $m$ . The margin is defined such that $\langle U_j, V_i \rangle \ge m$ if $A_{ji}=1$ (relevant) and $\langle U_j, V_i \rangle \le -m$ if $A_{ji}=0$ (irrelevant). A large margin is a proxy for representation quality, controlling robustness to perturbations, quantization error, and compositional generalization.

The core question is: What is the minimal dimension $d$ required to achieve the maximal possible margin $m_{rd}(+\infty, A)$ (the margin achievable in infinite dimensions) for a given relevance structure?

Methodology

The paper approaches the problem through a synthesis of communication complexity, compressed sensing, and learning theory.

Theoretical Framework: The authors analyze the "sign-rank" and "margin complexity" of the incidence matrix $A$ . They focus on the specific case where $A$ represents all possible $k$ -sparse rows (denoted $S_{n,k}$ ), corresponding to a scenario where every subset of $k$ documents has a unique query relevant exactly to that subset.
Upper Bound Construction (Theorem 1.4): To prove that low dimensions suffice, the authors employ the Johnson-Lindenstrauss (JL) Lemma. Unlike previous approaches that applied JL to all pairs of vectors, this work applies JL to the set of convex combinations of document embeddings. They utilize Maurey's Approximation Lemma to show that preserving the norms of these convex combinations is sufficient to guarantee the existence of query vectors with the desired margin in a reduced dimension.
Lower Bound Proof (Theorem 1.5): To establish tightness, the authors construct a large sphere packing argument. They demonstrate that if a margin exists, one can construct a set of vectors (linear combinations of document embeddings) that form a packing with a specific radius. By applying standard volume arguments to this packing, they derive a lower bound on the dimension required.
Alternative Constructions (Theorem 1.6): The paper explores constructions for dimensions smaller than the optimal $O(k \log n)$ regime. These rely on Khatri-Rao lifts (tensor products of random spherical vectors) and connections to the Restricted Isometry Property (RIP) to achieve non-trivial margins in dimensions as low as $O(k^2)$ .
Empirical Validation: The authors train embeddings in a "free embedding model" (optimizing vectors directly without a neural network architecture) using two loss functions: InfoNCE and Sigmoid. They compare the minimal dimension required to achieve a positive margin and the magnitude of the margin achieved.

Key Contributions and Results

1. Main Theoretical Result: Low-Dimensional Saturation

The paper's primary contribution is Theorem 1.4, which states that the optimal margin achievable in infinite dimensions can be nearly achieved in a dimension that grows only logarithmically with the data size.

Result: For any matrix $A$ , the dimension $d = O(m_{rd}(+\infty, A)^{-2} \log(\min(n, N)))$ is sufficient to achieve a margin of $(1-\epsilon)m_{rd}(+\infty, A)$ .
Implication for $k$ -sparse matrices: For the $S_{n,k}$ model, where the optimal margin is $\Theta(k^{-1/2})$ , the required dimension is $O(k \log n)$ . This significantly improves upon previous bounds (e.g., $O(k^2 \log n)$ from [BDES02]) and resolves the setup of [WBNL26].

2. Tight Lower Bounds

Theorem 1.5 provides a matching lower bound for the $S_{n,k}$ case.

Result: To achieve a margin $m = \Theta(k^{-1/2})$ , the dimension must be at least $\Omega(k \log(n/k))$ .
Significance: Combined with the upper bound, this proves that $d = \Theta(k \log(n/k))$ is necessary and sufficient for maximal margin retrieval in the $k$ -sparse setting. This fully resolves the theoretical limits for this specific model.

3. Constructions for Sub-Optimal Dimensions

Theorem 1.6 addresses the regime where $d = o(k \log(n/k))$ .

Result: Using Khatri-Rao lifts, the authors construct embeddings that achieve a margin of $\Omega(1/n)$ in dimension $d = \Theta(k^2)$ . This demonstrates a smooth trade-off between dimension and margin, showing that even in very low dimensions, non-trivial retrieval performance is theoretically possible.

4. Empirical Comparison of Loss Functions

The experiments reveal a stark difference between InfoNCE and Sigmoid losses in the free embedding model:

InfoNCE: Requires a dimension scaling with $n$ (consistent with [WBNL26]) to achieve a positive margin.
Sigmoid: Achieves a positive margin in a dimension close to the theoretical minimum (e.g., $d \approx 5$ for $k=2$ ) and largely independent of $n$ .
Explanation: The authors link this to Proposition 7, showing that the global minimizers of the sigmoid loss (with trainable inverse temperature) exactly coincide with the definition of margin embeddings, whereas InfoNCE does not share this property for $k > 1$ .

Significance and Claims

The paper claims to resolve a fundamental question in retrieval theory: Dimensionality is not a barrier for retrieval models.

Theoretical Resolution: The work establishes that high-quality retrieval (maximal margin) does not require high-dimensional representations. Instead, the dimension required scales logarithmically with the dataset size and linearly with the sparsity $k$ . This explains why practical models with $d \approx 1000$ can scale to billions of items.
Methodological Bridge: The paper successfully bridges the gap between communication complexity (sign rank), compressed sensing (RIP), and modern retrieval, providing new tools (Maurey's lemma application) for analyzing embedding quality.
Loss Function Insight: The empirical results suggest that the choice of loss function is critical for low-dimensional performance, with the sigmoid loss offering a theoretical and practical advantage for achieving large margins in compact spaces.

The authors remain modest regarding the scope, noting that their experiments are in a "free embedding" model and may not fully capture the complexities of real-world datasets or specific neural architectures. However, the theoretical bounds provide a rigorous foundation for understanding the geometric limits of retrieval.