Sinkhorn Based Associative Memory Retrieval Using Spherical Hellinger Kantorovich Dynamics¶

会议: NeurIPS 2026
arXiv: 2603.20656
OpenReview: https://openreview.net/forum?id=fgXY4eVgge
论文: https://arxiv.org/abs/2603.20656
领域: 优化/理论
关键词: 联想记忆, 经验测度, 德偏置Sinkhorn散度, 球面Hellinger-Kantorovich动力学, 粒子系统

一句话总结¶

本文将现代Hopfield联想记忆拓展至离散经验测度（加权点云），基于去偏Sinkhorn散度构建Log-Sum-Exp能量泛函，并导出显式协同更新支集位置与粒子权重的球面Hellinger-Kantorovich梯度流及其确定性检索算法，具备维度指数级的存储容量与几何收敛保证。

研究背景与动机¶

联想记忆作为认知计算与深度学习的重要基元，其核心机制在于通过内容寻址将受损或带噪的查询向量动态演化至表征记忆模式的局部能量极小值。近年来，基于Log-Sum-Exp（LSE）软极小能量构建的现代密集联想记忆（Modern Hopfield Networks）极大拓展了模型的存储容量，并在连续态设定下与Transformer架构中的注意力机制建立了优美的理论桥梁。然而，绝大多数现存联想记忆体系均将存储单元与查询局限于欧氏空间中的静态实向量，难以捕捉现实世界复杂模式中固有的不确定性、多模态演变以及种群层面的分布特征。

将联想记忆由向量空间拓展至概率测度空间是必然趋势，但现有测度记忆尝试存在不可忽视的理论与几何瓶颈。近期工作虽开始探索Bures-Wasserstein空间下的高斯分布联想记忆，但真实数据分布（如三维加权点云、经验后验粒子集、特征直方图）往往高度非高斯且非参数化。直接在离散经验测度上应用纯Wasserstein（Kantorovich）传输动力学会面临根本性阻碍：纯最优输运动力学仅移动粒子支集位置，而保持权重固定不变，完全丧失了调节粒子显著性、混合比例或自适应过滤噪声质点的能力。

本文从非平衡最优输运的内生几何出发，旨在构建一种能够同时演化质点位置与概率权重的离散测度密集联想记忆。核心 idea：基于去偏Sinkhorn散度构造测度空间上的Log-Sum-Exp能量泛函，并推导其在球面Hellinger-Kantorovich（SHK）几何下的连续梯度流与显式离散算法，实现支集重心的Kantorovich输运与单纯形权重的乘性反应机制协同更新。

方法详解¶

整体框架¶

算法整体处理流程包含三个主要层次：首先，将待存储模式与输入查询均建模为具有有限支集点且权重处于严格单纯形内部的经验测度；其次，利用可微且对称的去偏Sinkhorn散度构建平滑的Soft-min能量泛函，度量查询测度与各存储模式间的分布差异；最后，在球面Hellinger-Kantorovich（SHK）黎曼流形结构下沿能量负梯度方向离散演化，形成一个同时包含重心传输与复制动态重加权的确定性闭环检索系统。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入查询测度与存储模式<br/>加权点云粒子集 (a, x)"] --> B["熵正则最优输运求解<br/>计算Sinkhorn耦合矩阵与对偶势"]
    B --> C["去偏Sinkhorn散度与Gibbs权重分配<br/>能量泛函及其一阶变分计算"]
    C --> D["支集重心投影与输运更新<br/>Kantorovich位置移动"]
    C --> E["单纯形乘性重加权更新<br/>Spherical Hellinger反应动力学"]
    D --> F["粒子重组与下一轮迭代验证"]
    E --> F
    F -->|判定收敛或达到最大轮次| G["输出检索测度模式<br/>稳定收敛的吸引子 (a*, x*)"]

关键设计¶

1. 基于去偏Sinkhorn散度的测度Log-Sum-Exp能量泛函：解决离散测度平滑可微度量与自偏置问题 传统未正则化的Wasserstein距离在离散测度上缺乏处处光滑性，计算开销大且对偶势梯度奇异。本文利用熵正则化最优输运 $\operatorname{OT}_\varepsilon$ 定义无偏且处处平滑的去偏Sinkhorn散度： $$S_\varepsilon(\mu, \nu) \coloneqq \operatorname{OT}_\varepsilon(\mu, \nu) - \frac{1}{2}\operatorname{OT}_\varepsilon(\mu, \mu) - \frac{1}{2}\operatorname{OT}_\varepsilon(\nu, \nu)$$ 设存储模式为 $N$ 个离散经验测度 $\{X_i\}_{i=1}^N$，查询测度为 $\xi = \sum_{m=1}^M a_m \delta_{x_m}$。以此构建密集联想记忆的Soft-min能量泛函： $$E(\xi) = -\frac{1}{\beta}\log\left(\sum_{i=1}^N \exp\bigl(-\beta S_\varepsilon(\xi, X_i)\bigr)\right)$$ 其中 $\beta > 0$ 为逆温度系数。该能量泛函的中心化一阶变分 $u_\xi(x) \coloneqq \frac{\delta E}{\delta \xi}(\xi)(x) - \int_\Omega \frac{\delta E}{\delta \xi}(\xi) d\xi$ 能够通过各模式的Gibbs权重 $w_i(\xi) \propto \exp(-\beta S_\varepsilon(\xi, X_i))$ 及薛定谔对偶势 $f_{\xi, X_i}$ 显式闭式求出，完全消除了规范不变性（Gauge Invariance）带来的任意常数扰动。

2. 球面Hellinger-Kantorovich流形几何诱导的梯度分解：突破纯输运无法重加权粒子权重的瓶颈 纯Wasserstein动力学仅能推动支持点 $x_m$，但离散测度的粒子权重 $a_m$ 承载着质点显著性与浓度信息。为此，本文将离散测度流形 $\mathcal{M}_M = \Delta_{M, a_{\min}}^\circ \times \operatorname{Loc}_{M, \Delta_{\min}}(\Omega)$ 赋予球面Hellinger-Kantorovich（SHK）黎曼度量： $$\langle (\delta a, \delta x), (\delta a', \delta x') \rangle_{(a, x)} = \sum_{m=1}^M \left( \frac{\lambda^2}{a_m} \delta a_m \delta a'_m + a_m \delta x_m \cdot \delta x'_m \right)$$ 其中 $\lambda > 0$ 为控制反应（质量重加权）分量相对强度的尺度参数。通过变分与度量张量的对偶匹配，严格证明了能量 $E$ 的SHK黎曼梯度具有位置与权重的正交分解形式： $$\operatorname{grad}_{\operatorname{SHK}} E(a, x) = \left( \left(\frac{a_m}{\lambda^2} u_\xi(x_m)\right)_{m=1}^M, \, \left(\nabla u_\xi(x_m)\right)_{m=1}^M \right)$$ 这一设计使连续动力学退化为流体力学中的连续性-反应方程 $\partial_t \xi_t + \nabla \cdot (\xi_t v_t) = \xi_t r_t$，自然实现了空间漂移速度场与单纯形保质量反应速率的协同。

3. 对偶势梯度至重心映射的闭式投影：消除复杂ODE数值积分开销 在二次代价 $c(x,y)=\frac{1}{2}\|x-y\|^2$ 下，薛定谔势的梯度具有极为优雅的几何结构：$\nabla f_{\mu, \nu}(x) = x - T_{\mu \to \nu}^\varepsilon(x)$，其中 $T_{\mu \to \nu}^\varepsilon(x) = \int y \, d\pi_{\mu,\nu}^\varepsilon(y \mid x)$ 正是由最优熵正则输运耦合给出的重心投影映射（Barycentric Projection）。在离散粒子系统下，位置与权重的单步前向一阶Euler离散更新表达为纯确定性映射： $$x_m^{k+1} = x_m^k + \eta \left( \sum_{i=1}^N w_i^k T_i^k(x_m^k) - T_0^k(x_m^k) \right)$$ $$a_m^{k+1} \propto a_m^k \exp\left(-\frac{\eta}{\lambda^2} z_m^k\right), \quad \sum_{m=1}^M a_m^{k+1} = 1$$ 其中 $z_m^k = \sum_{i=1}^N w_i^k \left(f_i^k[m] - \frac{1}{2}(f_0^k[m] + g_0^k[m])\right)$。权重更新本质上构成了单纯形上的复制动态（Replicator Dynamics），不仅自然维持正值约束与单位归一化，更无需引入投影算子或松弛惩罚。

损失函数 / 训练策略¶

作为无需监督训练的密集联想记忆模型，记忆存储模式 $\{X_i\}_{i=1}^N$ 在初始化阶段直接固化，或通过采样机制 SampAlgo 构建。模型检索阶段由确定的离散迭代算子 $\Phi_\eta(\xi)$ 执行，其理论分析依赖以下关键性质： - 吸引盆不变性与局部PL条件：在Sinkhorn邻域 $B_i(r)$ 内，能量泛函满足局部收缩性 $\frac{1}{2}\|\operatorname{grad}_{\operatorname{SHK}} E(\xi)\|_{\operatorname{SHK},\xi}^2 \ge \mu (E(\xi) - E_i^*(r))$，保证误差按几何速率 $(1-\eta\mu)^{k/2}$ 线性收敛至唯一局部极小值 $X_i^*(r)$。 - 存储容量定理：利用高维球体上的Rademacher随机向量均值构造模式，证明当模式数量满足 $N = \lfloor \sqrt{2p} \exp(\frac{\gamma^2}{4}d) \rfloor$ 时，模式间的Sinkhorn吸引盆以不低于 $1-p$ 的概率两两不交，从而在环境维度 $d$ 上实现指数级存储容量。 - 吸引子保真度界：理论证明局部极小点与真实模式满足 $S_\varepsilon(X_i^*(r), X_i) \le \frac{N-1}{\beta} e^{-\beta\Delta}$，随着逆温度 $\beta$ 增大，检索吸引子指数逼近真实存储样本。

实验关键数据¶

主实验¶

论文在二维平面（$d=2$）的高斯点云经验分布记忆存储与恢复任务上，将提出的 SinkhornSHK Algo 与展平向量化的欧氏 Hopfield 模型（Euclidean Algo，将粒子坐标与对数权重拼接后利用标准内积与固定点迭代检索）进行对比。两组实验均设置 $N=5$ 个存储模式，逆温度 $\beta=50$，Sinkhorn 正则化参数 $\varepsilon=0.05$。

实验设定	模式数 $N$ / 点数 $M$	查询扰动噪声 $\sigma$	SinkhornSHK 成功检索率 (%)	Euclidean Hopfield 成功检索率 (%)	相对表现优势
实验1：高斯均值不同分布 ($\mu_i$ 分离)	$N=5, M=30$	$\sigma = 0.5$	100% (5/5)	100% (5/5)	两者均能利用全局均值位移准确恢复
实验2：零均值协方差差异分布 ($\mu_i = \mathbf{0}$)	$N=5, M=25$	$\sigma = 0.2$	100% (5/5)	40% (2/5)	SinkhornSHK 显著胜出（多恢复 3 例）

消融实验¶

针对动力学几何与参数敏感性（包括几何收敛速率界、逆温度与分离裕度对记忆保真度的影响）的理论与数值消融分析如下：

配置 / 机制	渐近收敛性界 / 误差上界	吸引盆重叠 / 失败模式分析	机制收益说明
完整 SHK 动力学 ($\eta < \eta_{\mathrm{ret}, i}$)	误差以 $(1-\eta\mu)^{k/2}$ 几何速率下降	吸引盆严格互斥保持 ($B_i(r) \cap B_j(r) = \emptyset$)	粒子位置与权重协同演化，保真度界达 $O(\frac{N}{\beta}e^{-\beta\Delta})$
仅纯 Kantorovich 输运 ($\lambda \to \infty$)	权重退化为冻结初值 $a^{(0)}$	当查询初始权重非均匀或带偏置时，收敛目标发生漂移	无法自适应修正显著度差异，限制在固定基数均匀测度
欧氏展平内积机制 (Euclidean Algo)	无分布不变性与输运度量保证	均值重合时几何测度信息丢失，在高斯噪声下陷入虚假吸引子	忽略排列不变性与测度内在传输代价，高混叠下失效

关键发现¶

在实验1中，各高斯模式具有明显分离的空间中心，欧氏向量化内积能够捕捉整体偏移，两者均实现 100% 检索；然而在实验2中，所有模式的均值完全重合为 $(0,0)$，仅依靠协方差结构与粒子高阶分布形状区分。欧氏展平距离完全忽视了点云粒子的置换不变性与最优输运几何，在 3 个模式实例中陷入混叠而检索失败，而 SinkhornSHK 算法通过熵正则对偶势精确度量分布相似性，全部 5 个案例均完美恢复。
动力学具有严格的几何收敛性与稳定性闭环：理论上由 PL 不等式驱动，实际仅需几十步重心迭代即可将查询推入真实模式的高保真邻域内；去偏修正项 $T_0^k(x_m^k)$ 对消除熵扩散引起的自收缩现象至关重要。

亮点与洞察¶

首次将密集联想记忆（Dense Associative Memory）拓展至离散经验测度空间，建立了一套严格基于非平衡最优输运理论（球面Hellinger-Kantorovich流形）的确定性迭代框架。
借助薛定谔对偶势梯度与重心投影映射之间的代数等价性，巧妙绕开了非平衡测度动力学中昂贵的速度场数值求解，直接以离散矩阵乘法形式给出了兼具显式位置漂移与单纯形复制动态的优雅更新律。
理论证明了该模型在高维空间下具有与经典向量现代Hopfield网络相媲美的指数级存储容量（$O(e^{c d})$），将联想记忆的容量保证严格延拓到了概率测度流形上。

局限与展望¶

论文数值实验目前仅在二维合成高斯点云上进行了概念验证（$N=5, d=2$），尚未在三维雷达点云（如 ModelNet40）、大规模分子构型或真实多模态直方图上进行高维扩展验证。
每一轮迭代均需针对 $N$ 个模式及查询自身求解熵最优输运（共 $N+1$ 次 Sinkhorn 算法），尽管具备完全确定性，但计算复杂度随存储模式数 $N$ 及粒子数 $M$ 线性/二次增长，在超大记忆库场景下亟待开发层次化剪枝或自适应缓存加速机制。
尚未考虑点云粒子数量不一致（$M$ 动态变化或带有非平衡总质量漂移）的情形，未来可结合无平衡非归一化最优输运进行泛化。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次建立非平衡最优输运（SHK）几何下的离散测度密集联想记忆，理论视角极具前瞻性。
实验充分度: ⭐⭐⭐☆☆ 数值实验仅限于二维玩具级高斯点云验证，缺乏大规模高维真实数据评测。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严丝合缝，几何直觉清晰，附录证明完备详尽。
价值: ⭐⭐⭐⭐☆ 为三维点云记忆寻址、粒子分布联想与基于能量的测度生成模型提供了扎实的数理基石。

实验设定	模式数 \(N\) / 点数 \(M\)	查询扰动噪声 \(\sigma\)	SinkhornSHK 成功检索率 (%)	Euclidean Hopfield 成功检索率 (%)	相对表现优势
实验1：高斯均值不同分布 (\(\mu_i\) 分离)	\(N=5, M=30\)	\(\sigma = 0.5\)	100% (5/5)	100% (5/5)	两者均能利用全局均值位移准确恢复
实验2：零均值协方差差异分布 (\(\mu_i = \mathbf{0}\))	\(N=5, M=25\)	\(\sigma = 0.2\)	100% (5/5)	40% (2/5)	SinkhornSHK 显著胜出（多恢复 3 例）

配置 / 机制	渐近收敛性界 / 误差上界	吸引盆重叠 / 失败模式分析	机制收益说明
完整 SHK 动力学 (\(\eta < \eta_{\mathrm{ret}, i}\))	误差以 \((1-\eta\mu)^{k/2}\) 几何速率下降	吸引盆严格互斥保持 (\(B_i(r) \cap B_j(r) = \emptyset\))	粒子位置与权重协同演化，保真度界达 \(O(\frac{N}{\beta}e^{-\beta\Delta})\)
仅纯 Kantorovich 输运 (\(\lambda \to \infty\))	权重退化为冻结初值 \(a^{(0)}\)	当查询初始权重非均匀或带偏置时，收敛目标发生漂移	无法自适应修正显著度差异，限制在固定基数均匀测度
欧氏展平内积机制 (Euclidean Algo)	无分布不变性与输运度量保证	均值重合时几何测度信息丢失，在高斯噪声下陷入虚假吸引子	忽略排列不变性与测度内在传输代价，高混叠下失效