Tan-HWG框架：用Wasserstein几何约束Hebbian学习实现稳健持续学习

1. 项目概述：当Hebbian学习遇见Wasserstein几何

最近在整理一些关于类脑计算和持续学习的笔记，发现一个挺有意思的框架，叫Tan-HWG。这个名字听起来有点唬人，但拆开来看，其实就是把两个看似不搭界的东西揉在了一起：一个是神经科学里经典的Hebbian学习规则，另一个是数学里用来度量概率分布距离的Wasserstein几何。这个框架的核心目标，是试图从计算的角度，去模拟和解释大脑中“记忆巩固”这个神奇又关键的过程。

简单来说，我们的大脑不是U盘，不是存进去就完事儿了。你今天学了一堆东西，睡一觉，有些记得更牢了，有些就模糊了。这个从短期记忆到长期记忆的转化、筛选和强化的过程，就是记忆巩固。传统的神经网络，尤其是处理连续学习任务时，经常面临一个叫“灾难性遗忘”的难题——学了新知识，就把旧的忘得一干二净。这显然和人脑的稳健学习能力相去甚远。

Tan-HWG框架的出发点就在这里。它认为，记忆的巩固不仅仅是突触权重的简单叠加或衰减，而应该被视为在神经连接构成的“参数空间”里的一种几何演化。Hebbian学习（“一起放电的神经元连接在一起”）提供了局部更新的动力，而Wasserstein距离则提供了一个全局的、分布层面的“尺子”，用来衡量和约束记忆表征在演化过程中的整体变化，防止它跑偏或者崩塌。这就像是在一个复杂的地形里规划路径，Hebbian规则告诉你下一步往哪个坡走（局部梯度），而Wasserstein几何则确保你整条路线不会偏离大方向太远，不会掉进沟里（全局稳定性）。

这个框架适合谁呢？如果你对神经科学启发的AI模型、持续学习、机器学习理论，特别是概率分布和几何方法在AI中的应用感兴趣，那么Tan-HWG提供了一个非常独特的视角。它不是一个即插即用的代码包，更像是一套理论工具和建模思想，能帮你重新思考如何让机器学得更像人——更稳健、更高效、更善于积累知识。

2. 核心原理拆解：Hebbian动力与Wasserstein约束如何协同

要理解Tan-HWG，我们得先把它名字里的两个核心部件拆开看明白，再看它们是怎么组装起来的。

2.1 Hebbian学习：不只是“一起放电”

Hebbian学习规则通常被简化为“fire together, wire together”。在计算神经科学和机器学习中，它有很多数学表达形式，最常见的一种是Oja规则，用于在线性神经元模型中学习主成分。其基本思想是，突触权重 ( w_{ij} ) 的更新量与突触前神经元 ( i ) 的活动 ( x_i ) 和突触后神经元 ( j ) 的活动 ( y_j ) 的乘积成正比：( \Delta w_{ij} \propto x_i y_j )。

但在Tan-HWG的语境下，我们需要更深入地理解它的角色：

• 局部性：Hebbian更新只依赖于当前时刻相连的两个神经元的局部活动信息。它没有全局视野，不知道整个网络的状态。
• 相关性驱动：它强化的是神经元活动之间的统计相关性。如果两个神经元经常同时活跃，它们之间的连接就会增强。
• 作为记忆形成的微观机制：在模型中，Hebbian规则驱动着短期记忆痕迹的形成。每一次经验（输入数据）都会在网络的连接权重上留下一个微弱的、局部的印记。

然而，纯粹的、不加约束的Hebbian学习是有问题的。它可能导致权重无限制增长（需要引入归一化机制），更重要的是，当面对连续输入的不同数据流时，局部的Hebbian更新很容易覆盖掉之前学到的模式，这就是灾难性遗忘的微观根源。它缺乏一个“宏观调控”机制来保护已经形成的、有价值的记忆结构。

2.2 Wasserstein几何：度量记忆的“形状”与“距离”

Wasserstein距离，也叫推土机距离（Earth Mover‘s Distance），是度量两个概率分布之间差异的一种方式。想象你有两堆土（两个分布），Wasserstein距离就是把一堆土挪动、重新塑造成另一堆土所需的最小“工作量”（成本）。

在Tan-HWG框架中，记忆不是被看作一个个孤立的权重数值，而是被表征为网络活动模式或连接权重的一种概率分布。例如，可以将网络中所有突触的权重值视为一个高维空间中的点云，这个点云的形状（分布）就编码了当前的记忆状态。

Wasserstein几何在这里扮演了几个关键角色：

1. 记忆状态的全局描述子：它不关心某个具体权重是0.5还是0.6，而是关心所有权重构成的整体统计形态（如均值、方差、高阶矩）。这种描述对噪声和微小扰动更鲁棒。
2. 记忆演化的约束器：这是核心。框架假设，健康的、稳健的记忆巩固过程，应该使得记忆表征的分布（即“记忆状态”）在Wasserstein度量下是平滑、连续演化的。也就是说，从时间 ( t ) 到 ( t+1 )，记忆分布的变化（Wasserstein距离）应该被限制在一个较小的范围内。
3. 定义记忆巩固的目标：记忆巩固可以被形式化为一个优化问题：在局部Hebbian动力的驱动下，寻找网络参数的演化路径，使得这条路径上任意两点记忆状态之间的Wasserstein距离积分（或某种代价）最小化。这相当于在Wasserstein几何定义的空间里，寻找一条“最短路径”或“最省力路径”来整合新记忆。

2.3 Tan-HWG的协同框架：微观动力与宏观几何的握手

现在，我们把两者结合起来。Tan-HWG框架的基本建模思路可以概括为以下几步：

1. 状态定义：将神经网络在时刻 ( t ) 的状态（通常是所有可训练参数，如权重和偏置）映射为一个概率分布 ( P_t )。这可以通过将参数向量视为样本，或通过其诱导的网络活动分布来实现。
2. 局部动力：当新数据 ( x_{t+1} ) 输入时，网络根据当前权重进行前向传播，产生活动。基于这些活动，应用某种形式的Hebbian学习规则（可能是经过修改的，包含衰减项）来计算权重的局部更新量 ( \Delta W_{local} )。
3. 全局几何约束：计算如果直接应用这个局部更新，新的参数分布 ( P_{t+1}^{local} ) 与旧分布 ( P_t ) 之间的Wasserstein距离 ( W(P_t, P_{t+1}^{local}) )。
4. 约束优化：框架引入一个正则化项或约束条件，要求 ( W(P_t, P_{t+1}) ) 不能超过某个阈值 ( \rho )。这里的 ( P_{t+1} ) 是实际要更新到的目标分布。因此，我们需要求解一个带约束的优化问题：在 ( W(P_t, P_{t+1}) \leq \rho ) 的条件下，让更新后的网络性能（如对新数据的拟合程度）尽可能好，同时尊重Hebbian动力指出的方向。
5. 记忆巩固的实现：上述优化问题的解，就是最终执行的权重更新 ( \Delta W_{final} )。这个过程可以理解为：Hebbian规则提出了一个“原始草案”，而Wasserstein几何约束则对这个草案进行“审议和修正”，确保修改后的方案不会破坏已有的整体记忆结构。这个“审议修正”的过程，在计算上就对应了记忆的巩固——新经验被整合进来，但整合的方式是受控的、平滑的，避免了对旧记忆的粗暴覆盖。

注意：这里的“Tan”可能指代“Tangent”（切空间）。在实际的数学处理中，由于直接在高维分布空间优化Wasserstein距离计算量巨大，一个常见的技巧是利用Wasserstein空间的几何结构，在其切空间（Tangent Space）进行局部近似和优化。这可能是框架名称中“Tan”的由来，即将全局的Wasserstein约束，转化为每次更新时在局部切空间上的线性约束，使得算法可行。

3. 算法实现的关键步骤与数学形式化

理解了核心思想后，我们来看看如何将其转化为一个可计算、至少是可形式化的算法框架。这里会涉及一些数学，但我会尽量用直观的方式解释。

3.1 将网络参数分布化

第一步是如何将确定的网络参数 ( \theta ) 变成一个分布 ( P )。有两种主流思路：

方法一：参数空间的经验分布将网络的 ( N ) 个参数（例如所有权重）扁平化成一个高维向量 ( \theta \in \mathbb{R}^N )。我们可以简单地将其视为 ( N ) 维空间中的一个点。为了形成分布，一个实用的技巧是引入一个微小的随机扰动。例如，考虑参数 ( \theta ) 的一个小邻域，或者在进行优化时，使用 mini-batch 梯度下降带来的隐式随机性。更形式化地，我们可以定义 ( P_\theta ) 为以 ( \theta ) 为均值的某个简单分布（如各向同性高斯分布），尽管在计算距离时，我们通常关注的是由 ( \theta ) 直接定义的“点质量”分布与另一个点分布之间的距离。

方法二：神经元活动空间的分布这是更贴近神经科学解释的一种方式。给定一个固定的输入分布（或一组锚定输入），网络每一层神经元的激活值构成一个分布。记忆可以被认为是这些激活模式之间的关联，因此用激活值的分布来表征记忆状态更为直接。假设我们有一组参考输入 ( {x^{(i)}}{i=1}^M )，通过网络前向传播，在某一层（通常是隐藏层）得到激活向量集合 ( {h^{(i)}}{i=1}^M )。这个集合就定义了该层在当前网络参数 ( \theta ) 下的激活经验分布 ( P_\theta^h )。

在Tan-HWG中，第二种方法可能更具解释力，因为Hebbian学习直接作用于神经元活动。但第一种方法在数学处理上可能更简单。实际框架可能会根据具体问题灵活选择。

3.2 Wasserstein距离的计算与近似

精确计算两个高维经验分布之间的Wasserstein距离（特别是2-Wasserstein距离 ( W_2 ) ）是计算昂贵的，复杂度可达 ( O(n^3 \log n) )。在机器学习中，我们不得不使用近似。

1. 切片Wasserstein距离：一种高效且常用的近似是切片Wasserstein距离。其思想是：不是直接在高维空间计算距离，而是随机抽取许多个一维方向（切片），将高维分布投影到这些方向上，计算所有一维投影分布之间的Wasserstein距离的平均值。数学上，对于单位球面上的随机方向 ( \gamma )，有： ( SW(P, Q) = \mathbb{E}{\gamma \sim U(S^{d-1})} [W_1(P\gamma, Q_\gamma)] ) 其中 ( P_\gamma, Q_\gamma ) 是分布沿方向 ( \gamma ) 的投影。这极大地降低了计算成本，且在实践中被证明是有效的。

2. 基于最优传输正则化的代理损失：另一种思路是不直接计算距离，而是将其作为优化问题的一个正则化项。例如，在更新参数从 ( \theta_t ) 到 ( \theta_{t+1} ) 时，我们希望最小化损失函数 ( L(\theta_{t+1}) )（如预测误差），同时约束 ( W_2(P_{\theta_t}, P_{\theta_{t+1}}) \leq \rho )。利用Wasserstein距离的对偶形式，这个约束可以转化为在损失函数中添加一个惩罚项，其梯度与 ( \theta_t ) 和 ( \theta_{t+1} ) 处损失函数的梯度差有关。这引出了一类称为“Wasserstein梯度惩罚”的方法。

在Tan-HWG的上下文中，算法很可能采用一种在线近似。每次接收到新数据 mini-batch 时：

• 用当前参数 ( \theta_t ) 计算网络在 batch 上的激活分布 ( P_t )。
• 执行一步基于Hebbian的局部更新（可能是梯度下降的一个变体），得到临时参数 ( \theta_{temp} ) 及其激活分布 ( P_{temp} )。
• 快速估算 ( W(P_t, P_{temp}) )（例如用切片法）。
• 如果距离超过阈值 ( \rho )，则对更新方向进行修正。修正的方法可以是在梯度方向上回退一步，或者将更新向量投影到满足Wasserstein约束的可行子空间（这涉及到在Wasserstein空间的切空间上进行操作，即“Tan”的部分）。

3.3 Hebbian更新与Wasserstein约束的融合算法

一个简化的算法伪代码流程可能如下：

初始化网络参数 θ， 阈值 ρ， 学习率 η。 对于每一个训练批次（或每一个新数据点）： 1. 前向传播：用当前θ处理批次数据，得到各层激活。 2. 计算局部Hebbian更新 Δθ_local： - 根据激活计算相关性矩阵。 - 根据选定的Hebbian规则（如Oja, BCM等）计算权重变化。 - Δθ_local = Hebbian_Update(activations)。 3. 计算临时参数：θ_temp = θ + η * Δθ_local。 4. 评估分布变化： - 基于当前θ和θ_temp，分别计算在某个参考数据集或当前批次上的激活分布 P 和 P_temp。 - 近似计算 Wasserstein 距离 d = SW(P, P_temp)。 5. 应用约束更新： if d <= ρ: θ = θ_temp # 接受Hebbian更新 else: # 需要进行约束修正 # 方案A：缩放更新步长 scale_factor = ρ / d θ = θ + scale_factor * η * Δθ_local # 方案B：求解带约束的优化问题（更精确但更慢） # θ = argmin L(θ') s.t. W(P_θ, P_θ') <= ρ, 其中L是结合了任务损失和Hebbian驱动的目标函数。 6. （可选）执行任务特定损失的精调：在Wasserstein约束的大框架下，可能还会用一个小的标准梯度下降步来微调，以优化主要任务（如分类）的损失。

这个流程体现了“提议-修正”的思想。Hebbian规则是激进的“提议者”，总想根据最新经验改变连接。Wasserstein约束是稳健的“守护者”，确保任何改变都不会让记忆的整体形态发生剧变。两者的平衡点，就是记忆巩固发生的地方。

实操心得：阈值 ( \rho ) 是一个超参数，它控制着记忆的“可塑性-稳定性权衡”。ρ 太大，约束太弱，可能无法有效防止遗忘；ρ 太小，约束太强，网络会变得僵化，难以学习新知识。在实际调参中，可能需要设置一个随时间衰减的 ρ，模拟学习初期可塑性强、后期稳定性增强的过程。

4. 应用场景与优势分析

Tan-HWG不是一个通用网络架构，而是一个学习范式或元框架。它可以被应用到多种需要持续学习和记忆巩固的场景中。

4.1 持续学习与灾难性遗忘缓解

这是最直接的应用。在经典的持续学习基准测试上（如Split MNIST, Permuted MNIST, Incremental Class Learning），我们可以将Tan-HWG框架应用于多层感知机或卷积神经网络。

• 如何应用：将网络每一层的权重更新，不仅由任务损失（如交叉熵）的梯度决定，还同时受到该层神经元活动分布Wasserstein距离的约束。Hebbian成分可以体现为一种基于激活相关性的内在驱动，与任务梯度结合。
• 预期优势：相比于传统的弹性权重巩固或生成回放等方法，Tan-HWG从分布几何层面进行约束，可能提供更平滑、更理论可解释的遗忘控制。它不依赖于保存旧数据或精确计算参数的重要性，而是通过维护激活模式的整体稳定性来保护旧技能。

4.2 脉冲神经网络中的生物可塑性模拟

SNN天然使用脉冲时间依赖可塑性等类Hebbian规则。Tan-HWG可以为SNN提供一套理论框架，来模拟更复杂的生物记忆巩固现象，如睡眠期间的记忆重播和突触缩合。

• 如何应用：将STDP视为微观的Hebbian动力。在模拟睡眠或安静清醒期时，让网络在无外部输入的情况下运行，内部活动由噪声或自反馈驱动。此时，应用Wasserstein约束下的“虚拟”Hebbian更新，可以让网络自发地重新组织连接，强化重要的活动模式（记忆重播），削弱不重要的连接（突触缩合），从而在分布层面优化记忆表征。

4.3 强化学习中的技能巩固与迁移

在强化学习中，智能体需要在一系列任务中积累技能。Tan-HWG可以用于约束策略网络或价值网络的更新。

• 如何应用：当智能体在新环境中学习时，其策略网络参数的变化，会受到旧策略下状态-动作访问分布与新策略下分布之间Wasserstein距离的限制。这可以防止新学习完全覆盖旧策略，有利于技能的积累和向新任务的迁移。Hebbian成分可以对应于基于TD误差的某种神经可塑性模拟。

4.4 框架的优势与潜在挑战

优势：

1. 统一的几何视角：它将记忆巩固抽象为一个在表征分布空间中的几何流形学习问题，提供了强大的数学工具进行分析。
2. 生物合理性：同时考虑了局部突触规则（Hebbian）和全局系统级约束（分布稳定性），与神经科学中多尺度理解大脑的观点相符。
3. 理论优雅性：Wasserstein距离具有良好的数学性质（度量性质、对弱收敛敏感），使得理论分析（如收敛性、稳定性）成为可能。
4. 避免显式数据回放：与基于回放的方法不同，它通过几何约束来保护记忆，可能更节省内存。

挑战与注意事项：

1. 计算开销：即使使用切片法等近似，计算和优化Wasserstein距离仍然比标准深度学习昂贵得多。这限制了其在大规模网络和数据集上的直接应用。
2. 分布定义的选择：如何定义“记忆分布” ( P ) 非常关键且非平凡。是基于参数、基于激活、还是基于网络输出的预测分布？不同的选择会导致完全不同的算法行为和解释。
3. 超参数调优：阈值 ( \rho )、学习率、Hebbian更新强度等需要精心调整。特别是 ( \rho )，它直接决定了稳定性和可塑性的平衡点。
4. 与深度学习的结合：如何将这一框架无缝集成到基于反向传播的深度网络训练中，需要巧妙的工程设计和近似。例如，如何计算Wasserstein距离相对于网络参数的梯度？

5. 实现难点与常见问题排查

在实际尝试实现或复现Tan-HWG思想时，你会遇到几个典型的坑。这里结合我的理解，分享一些排查思路。

5.1 分布漂移与阈值ρ的设定

问题描述：网络完全学不进去新任务，或者相反，学新任务时旧任务性能暴跌。这通常是因为Wasserstein距离阈值 ( \rho ) 设置不当。

• 排查与解决：
  1. 监控距离：在训练过程中，实时绘制每个任务切换前后、分布距离 ( d ) 的变化曲线。观察 ( d ) 的典型值范围。
  2. 动态调整：不要使用固定 ( \rho )。可以尝试：
     • 基于百分位：将 ( \rho ) 设置为历史 ( d ) 值的某个百分位数（如90%）。
     • 任务自适应：为每个已学任务维护一个基线分布 ( P_{task} )。学习新任务时，约束是针对所有旧任务基线分布中最严格的那个（即最小的允许距离）。
     • 渐进收紧：训练初期用较大的 ( \rho )（高可塑性），随着学习任务增多，逐渐减小 ( \rho )（增强稳定性）。
  3. 校准实验：在一个简单的双任务序列上，系统地扫描 ( \rho ) 值，观察新旧任务准确率随 ( \rho ) 变化的权衡曲线，找到“拐点”。

5.2 Wasserstein距离估计的不稳定性

问题描述：切片Wasserstein距离的估计方差较大，导致约束时紧时松，训练过程震荡。

• 排查与解决：
  1. 增加切片数：这是最直接的方法。虽然会增加计算量，但能有效降低估计方差。可以从几十个切片开始，逐步增加直到训练曲线平滑。
  2. 使用正则化最优传输：考虑使用Sinkhorn距离（熵正则化的Wasserstein距离）作为替代。它可以通过迭代矩阵缩放快速计算，且更平滑、更稳定。虽然它不再是严格的度量，但对于约束目的而言通常足够。
  3. 滑动平均：对计算出的距离 ( d ) 进行滑动平均（指数移动平均），用平滑后的值去判断是否违反约束，可以过滤掉高频噪声。
  4. 检查分布采样：确保用于计算距离的激活样本是独立同分布的，并且数量足够。如果batch size太小，分布估计本身就不准。

5.3 Hebbian更新与任务梯度的冲突

问题描述：纯粹的Hebbian更新方向可能与最小化任务损失（如分类误差）的梯度方向不一致甚至相反，导致学习效率低下。

• 排查与解决：
  1. 混合更新：不要完全用Hebbian更新代替梯度下降。采用加权和：( \Delta \theta = \alpha * \Delta \theta_{task} + \beta * \Delta \theta_{hebbian} )。其中 ( \Delta \theta_{task} ) 是标准反向传播梯度。通过调整 ( \alpha ) 和 ( \beta ) 来平衡。
  2. 任务驱动的Hebbian：修改Hebbian规则，使其受任务信号调制。例如，只对与正确分类相关的神经元活动进行Hebbian增强（类似于奖励调制的STDP）。
  3. 投影法：先计算任务梯度 ( g_{task} )，然后计算Hebbian更新 ( g_{hebb} )。将 ( g_{hebb} ) 投影到与 ( g_{task} ) 不正交的子空间上，或者只采用 ( g_{hebb} ) 中与 ( g_{task} ) 方向一致的分量。
  4. 分阶段训练：在一个mini-batch内，先使用任务梯度进行主要更新，然后施加一个很小的、受Wasserstein约束的Hebbian更新作为“微调”，模拟巩固过程。

5.4 计算效率与可扩展性

问题描述：算法太慢，无法用于稍大一点的网络或数据集。

• 排查与解决：
  1. 局部约束：不必对全网所有层的联合分布进行约束。可以对每一层单独施加Wasserstein约束，这会大大降低分布维度和计算量。神经科学上也有证据表明，记忆巩固可能在不同脑区局部发生。
  2. 低维投影：在计算距离前，先用PCA或自编码器将高维激活降维到一个有意义的低维子空间，然后在低维空间计算Wasserstein距离。
  3. 近似距离函数：训练一个小的神经网络（批评家）来近似两个激活分布之间的Wasserstein距离。一旦训练好，在前向传播中直接调用这个批评家网络来快速估计距离，代替昂贵的切片计算。
  4. 稀疏化：只对网络中一部分关键连接（如输出层或高层特征层）施加强约束，其他层使用较弱的约束或传统正则化。

实现Tan-HWG这类前沿框架，最大的收获往往不是得到一个SOTA的模型，而是在尝试解决上述问题的过程中，对“学习”、“记忆”和“表示”这些根本概念产生更深刻的理解。它迫使你从权重数值的层面跳出来，去思考网络整体表征空间的几何形状是如何随着学习而演化的。这种视角的转换，本身就是一个非常有价值的思维训练。