🤖 machine learning

Discovering Data Structures: Nearest Neighbor Search and Beyond

本文提出了一种通用的端到端学习框架，该框架能够在无需初始化的条件下，自动从零开始发现最优的数据结构和查询算法，成功复现了诸如二分查找、k-d树以及用于最近邻搜索的局部敏感哈希等已知方案，同时还能适应数据流中的频率估计。

原作者： Omar Salemohamed, Laurent Charlin, Shivam Garg, Vatsal Sharan, Gregory Valiant

发布于 2026-06-09

📖 1 分钟阅读☕ 轻松阅读

原作者： Omar Salemohamed, Laurent Charlin, Shivam Garg, Vatsal Sharan, Gregory Valiant

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你拥有一个巨大且杂乱无章的书库。传统上，图书管理员（计算机科学家）会花费数年时间设计特定的规则和归档系统（数据结构）。他们可能会说：“把所有书按字母顺序排列在书架上”或者“按颜色和大小进行分组”。这些规则对每个人都适用，但它们并不了解你的特定习惯。也许你总是借阅推理小说，或者你的图书馆有一个奇怪的模式，即 90% 的书都是关于猫的。

这篇论文提出了一个大胆的问题：我们能否通过观察书籍并练习如何寻找它们，教计算机从零开始发明自己的图书归档系统？

作者们说：可以。他们创建了一个“学习机器”，它不仅仅是遵循规则，而是能够自行发现规则。

双人组队

他们构建的系统就像是一个由两个机器人组成的协作团队：

组织者（数据处理网络）： 这个机器人观察杂乱的数据堆（书籍），并找出重新排列它们的最佳方式。它不仅仅是按字母顺序排序；它学习如何以一种能让下一个机器人工作更轻松的方式来进行排序。
搜索者（查询执行网络）： 这个机器人被赋予一个特定的问题（例如，“寻找关于猫的书”）。它只能窥视极少量的书架（受限的“查看预算”）。它必须学习一种策略，以便利用这仅有的几次窥视尽可能快地找到正确的书。

神奇之处在于它们协同训练。组织者学习如何专门为了帮助搜索者而排列书籍，而搜索者则学习如何阅读组织者的排列方式。它们练习了数百万次，直到发明出一套适用于该特定类型书籍的完美系统。

他们发现了什么？

研究人员在不同的“图书馆”（数据集）上测试了该系统，发现机器人重新发明了著名的各种人类发明，且往往优于这些发明：

简单列表（一维数据）： 当数据仅仅是一行数字时，组织者学会了完美地排序这些数字。随后，搜索者学会了一种比标准“二分查找”（类似于猜测列表中间位置）更好的策略。如果数字通常较小，搜索者就会学会从列表的开头开始查找，而不是从中间开始，从而节省时间。
二维地图： 当数据具有两个维度（如带有 X 和 Y 坐标的地图）时，机器人学会了构建 k-d 树。这是一种将地图分割成越来越小的方块以快速定位的复杂方法。机器人是在没有人告诉它们什么是“树”或“分割”的情况下，自行悟出了这一点。
高维迷宫： 在处理像图像这样具有数千个特征的复杂数据时，机器人学会了所谓的局部敏感哈希（LSH）。想象一下，看到一张猫的照片，就能立刻知道它属于“猫桶”，而无需查看其他所有照片。机器人学会了将复杂的图像投影到简单的桶中，就像人类专家所做的那样。
“热门项”技巧： 在一个涉及计数项出现频率的测试中（例如追踪互联网上频繁出现的 IP 地址），机器人学会了在内存中为最频繁出现的项目保留特殊的“VIP 插槽”。这防止了常见项目与稀有项目混淆，击败了标准的计数工具。

“顿悟”时刻

最令人惊讶的部分是，机器人不需要人类说：“嘿，试着排序一下！”或者“使用树状结构！”它们从随机噪声开始，通过试错，逆向工程出了这些经典的计算机科学算法。

在一次关于数字图像的实验中，机器人学会了识别这些图像实际上是数字，按数值进行排序，然后高效地进行搜索——这一切都没有被告知什么是“数字”或如何排序。它们只是学会了“将相似的图像归为一类”可以使搜索变得更快。

局限性（代价）

论文诚实地说明了其局限性：

规模： 实验是在相对较小的“图书馆”（约 100 到 500 个项目）上进行的。现实世界的图书馆拥有数百万个项目。目前的机器人可能会被如此大量的数据压垮。
速度： 机器人在开始搜索之前需要很长时间进行“思考”（预处理）。在现实生活中，我们通常需要即时答案。
黑盒： 虽然机器人找到了优秀的解决方案，但我们并不总能拥有一个简单的数学证明来解释为什么它们的特定排列方式有效。我们只知道它有效，因为我们测试过它。

核心结论

这篇论文证明了神经网络可以充当算法发明家。与其由人类设计归档系统，我们可以让计算机根据它所看到的特定数据模式，去发现最高效的组织和搜索数据的方法。这就像是给一个机器人一个乱糟糟的房间和有限的时间去寻找特定的玩具，然后看着它发明出一种甚至比人类设计的还要好的整理房间的方法。

技术摘要：端到端发现数据结构

问题陈述

本文研究了神经网络是否能够在没有人工设计先验或精心初始化的情况下，从零开始发现数据结构和查询算法。研究重点是最近邻（NN）搜索问题，即给定数据集 $D$ 和查询 $q$ ，目标是找到 $y = \arg\min_{x \in D} \text{dist}(x, q)$ 。挑战在于如何将这一组合优化问题构建为一个端到端的学习任务，同时对空间复杂度（内存大小）和查询复杂度（查找次数）施加约束，并确保所学到的结构对于特定的数据分布是高效的。

研究中探讨的第二个应用场景是流式环境下的频率估计，即模型在内存受限的情况下，观察元素流并估计某个查询项的频率。

方法论：发现框架

作者提出了一个由两个组件组成的端到端训练神经架构：

数据处理网络：
- 角色： 将原始数据集 $D$ 转换为结构化数据库 $\hat{D}$ 。
- 架构： 主要采用基于 NanoGPT 的 8 层 Transformer（具有二次方注意力机制），尽管文中也探索了用于提高效率的线性注意力变体。
- 机制： 该网络并非直接输出转换后的数据，而是为每个数据点学习一个排名（标量值）。随后通过一个可微排序函数根据这些排名对数据集进行重排序，从而生成 $\hat{D}$ 。这种基于排列的方法避免了学习重现原始数据值的困难，同时保留了表示任何顺序的能力。
- 扩展： 在允许额外空间的情况下，该网络可以输出额外的预计算向量或统计量，以辅助查询阶段。
查询执行网络：
- 角色： 通过对 $\hat{D}$ 进行固定预算 $M$ 次的查找来检索查询 $q$ 的答案。
- 架构： 由 $M$ 个独立的 MLP 序列 $Q_1, \dots, Q_M$ 组成。
- 机制： 每个查询模型 $Q_i$ 接收查询 $q$ 以及之前查找的历史记录 $H_i = \{(m_1, v_1), \dots, (m_{i-1}, v_{i-1})\}$ 作为输入。它输出一个 one-hot 向量 $m_i$ ，指示应访问 $\hat{D}$ 中的位置。
- 可微性： 为了处理训练过程中内存访问的离散性质，模型使用带有噪声注入的 Gumbel-Softmax。这鼓励了稀疏性和鲁棒性，使得模型能够通过梯度下降学习离散的选择策略。在推理阶段，则应用 hardmax。

训练目标： 系统在三元组 $(D, q, y)$ 上进行训练，其中 $y$ 是真实的最近邻。通过最小化预测输出 $\hat{y}$ 与 $y$ 之间的损失（根据设置使用平方损失或交叉熵）来进行优化。整个训练过程是端到端的，同时优化数据结构的构建和查询算法。

核心贡献与结果

1. 恢复经典算法

该框架在没有显式监督的情况下，成功发现了已知的最优算法：

一维最近邻搜索：
- 在均匀分布上，模型学习对数据进行排序，并执行一种优于二分查找的搜索策略，通过利用分布知识，使其能更接近目标值，类似于插值搜索。
- 在**“困难”分布**（即查询位置不提供先验信息时）上，模型发现了二分查找，达到了最坏情况下的最优性。
- 在 Zipf 分布下，模型能够适应偏斜分布，表现优于标准的二分查找。
二维最近邻搜索：
- 在均匀分布上，模型学习了一种将空间上接近的点聚集在一起的排列方式，其表现优于标准的 k-d 树。
- 在困难分布上，模型学习了一种递归划分策略，其行为与 k-d 树高度相似，通过在交替维度上进行数据分割。

2. 高维表示与 LSH

超球面数据： 在 30 维单位超球面上，模型学习到了一种类似于**局部敏感哈希（LSH）**的结构。它发现了一个投影到两个向量上的方案，从而对空间进行划分，有效地将相似的点映射到同一个“桶”中，实现高效检索。
MNIST（三位数）： 在被要求基于数值标签（而非像素距离）寻找最近邻时，模型学会了：
1. 从图像中提取捕捉数值顺序的一维特征。
2. 对这些特征进行排序。
3. 执行类似插值的搜索。
  这一切都是在没有任何关于数值标签或度量空间的显式监督下，完全从零开始实现的。
现实世界基准测试： 在 FashionMNIST 和 SIFT 数据集上，模型的表现与人工设计的 Learning-to-Hash 基准模型（如 ITQ、K-Means、NeuralLSH）相当。

3. 流中的频率估计

应用于流式频率估计问题（Zipf 分布）时，该模型优于标准的 CountMinSketch (CMS) 算法。
洞察迁移： 学习到的模型使用了比标准 CMS（使用 $\Delta=1$ ）更小的更新增量（increment value）。作者利用这一洞察设计了一个增强型 CMS，采用了自定义的更新增量，这使得它在处理真实世界的 CAIDA IP 流量数据集时，性能提升了高达 $1.9\times$ 。

4. 额外空间的使用

该框架展示了学习如何使用额外内存的能力。在一维和 30 维设置中，数据处理网络学习了预计算并存储统计量（例如，排序输入的加权和），这使得查询网络能以更少的查找次数找到最近邻。

重要意义与主张

本文声称迈出了端到端数据结构学习的第一步。其主要意义在于：

可行性： 证明了深度学习模型可以从零开始发现有意义的经典数据结构（排序、二分查找、k-d 树、LSH）和查询算法。
分布适应性： 表明学习到的结构可以适应特定的数据分布（例如，均匀分布 vs 困难分布 vs Zipf 分布），从而超越针对最坏情况设计的经典结构。
统一框架： 提出了一个将数据结构问题构建为端到端学习任务的通用范式，该范式不仅适用于最近邻搜索，还适用于频率估计、图结构和稀疏矩阵压缩等问题。
洞察生成： 强调了小规模学习模型可以产生可迁移的洞察，从而改进经典算法（如增强型 CountMinSketch 的案例）。

局限性与范围

作者明确指出了几点局限性：

规模： 实验仅限于小型数据集（ $N \approx 100$ ，最高至 $N=500$ ）。由于计算成本（二次方注意力和大规模输出向量的内存消耗）较高，将该方法扩展到实际的大规模数据集仍是一个挑战。
效率： 研究侧重于空间和查询复杂度。它并未充分优化预处理时间或神经网络的推理成本，而这些对于实际部署至关重要。
保证： 与提供可证明保证的学习增强算法不同，端到端学习的结构缺乏理论边界和可解释性。
失效模式： 文中指出，在特定设置下（例如具有随机排名顺序的高偏斜 Zipf 分布），学习可能会失败，此时可能需要精心的初始化或归纳偏置。