PARSING NEURAL DYNAMICS WITH INFINITE RECUR-RENT SWITCHING LINEAR DYNAMICAL SYSTEMS

无限递归切换线性动态系统

https://openreview.net/pdf?id=YIls9HEa52

摘要

神经活动和行为的无监督降维方法为理解神经信息处理的基础提供了前所未有的见解。一种流行的方法涉及循环切换线性动态系统（rSLDS）模型，该模型使用有限数量的低维线性动态系统之间的离散切换来描述神经尖峰序列数据的潜在动态。然而，rSLDS 模型的一些特性限制了其在试验变化数据上的部署，例如试验中固定数量的状态，以及缺乏状态的潜在结构或组织。在这里，我们通过为 rSLDS 模型赋予具有潜在几何结构的半马尔可夫离散状态过程来克服这些限制，该过程捕获了具有灵活状态基数的分区上随机过程的关键属性。我们利用偏微分方程（PDE）理论为离散状态的动态充分统计量推导出一种高效的半参数公式。这一过程与切换动态相结合，定义了我们的无限循环切换线性动态系统（irSLDS）模型类。我们首先在合成数据上验证并展示了我们模型的能力。接下来，我们转向分析小鼠在决策过程中的电生理数据，并揭示了试验内和试验平均神经活动背后的强非平稳过程。

1 引言

将行为和神经动态解析为更简单的交替状态，为理解大脑计算的本质提供了前所未有的见解（Wiltschko 等人，2015；Calhoun 等人，2019；Ashwood 等人，2022；Bolkan 等人，2022）。这类通常在多个试验中收集且跨越相当长时间的数据，往往与在多个时间尺度上表现出复杂动态的潜在过程相关（Cowley 等人，2020；Roy 等人，2021；Nassar 等人，2019；Yu 等人，2009）。学习就是这种非平稳性的一个重要例子，但也可能存在其他变化，如疲劳（Marcora 等人，2009）或唤醒（Schriver 等人，2018）。然而，从数据中推理这些潜在过程，无论是在试验内还是跨试验，仍然具有挑战性。

执行这种离散状态分割的状态空间模型通常围绕隐马尔可夫模型（HMM）展开。该模型假定存在一个潜在的、隐藏的离散马尔可夫链，每个状态有不同的观察模型来生成数据。每个离散状态还可以扩展以捕获不同的动态机制，控制某些连续空间随机过程的动态。由此产生的切换状态空间模型捕获了在离散动态机制集之间交替的活动，并已证明在复杂非线性活动的建模中有用（Fox 等人，2010；Smith 等人，2021）。在这类模型中捕获在线依赖关系的一个重要步骤是通过循环（Linderman 等人，2016；Zoltowski 等人，2020），其中连续动态变量控制离散状态之间的切换。这类模型为理解这些过程的时间依赖性提供了强大的见解（例如 Glaser 等人（2020）），但由于离散状态的固定基数和缺乏状态的几何结构，在训练时面临挑战。

事实上，传统的 HMM 以及基于它们的模型（如切换状态空间模型）使用固定数量的离散状态。这种固定基数通常通过交叉验证确定，这一过程可能证明计算成本很高。此外，它不鼓励在任何试验或会话子集上使用更少的离散状态，因此它们无法根据需要进行解释和确定。为了克服这些限制，多个途径考虑为 HMM 扩展灵活的状态基数（Beal 等人，2001；Fox 等人，2011）。这些模型围绕分层狄利克雷过程（HDP）展开，该过程在（理论上无限的）状态之间的转移概率上放置了灵活的非参数先验。不幸的是，尽管 HDP 及其推广可以捕获状态依赖性和持续性（Teh 等人，2006；Fox 等人，2008），但在其标准公式中，它们不允许离散状态分配之间的循环和其他依赖关系。

此外，HMM 没有任何先验的状态几何结构。特别是，离散状态是排列不变的，因此模型在离散状态标签顺序的任何排列下都是等价的（见附录图 4B 中的示意图）。这导致不可识别性，并限制了推理出的潜在离散状态的可解释性。尽管最近的工作集中在离散动态机制的可变性上（Nassar 等人，2019；Linderman 等人，2019；Costacurta 等人，2022），在灵活性方面取得了令人印象深刻的结果，但其中大多数可以被视为 “局部” 变化；它们考虑与每个离散状态相关的参数的扰动或参数化变化。据我们所知，没有专门解决离散状态空间的几何性质的。

在这项工作中，我们使用偏微分方程（PDE）理论为循环切换状态空间模型中的离散状态转移开发了一种替代先验。我们的公式使我们能够：（1）捕获具有灵活状态基数的随机过程的重要属性，同时支持循环连接的使用；（2）在离散状态上诱导几何结构；（3）诱导半马尔可夫离散状态动态以捕获更复杂的时间依赖性。

在第 2 节回顾相关背景之后，我们在第 3 节介绍我们的模型，无限循环切换线性动态系统（irSLDS），并提供其生成形式和推理过程的详细信息。在第 4 节中，我们首先在合成数据上验证模型并展示其特性，然后转向来自国际脑实验室（IBL）的小鼠电生理数据。PDE 先验由比传统基于 HMM 的转移模型更少的参数定义，但我们表明，它在保持甚至优于传统模型的同时，提供了可解释的结构，适用于揭示试验变化的结构。

2 背景

2.1 切换线性动态系统

3 无限循环切换线性动态系统

我们的技术贡献在于利用分布式中国餐馆过程（dist-CRP）的关键特性来引导 rSLDS 中的切换过程，从而支持高效的推理和生成。第 3.1 节介绍了简单组合的缺陷，详细说明了在这种情况下依赖波利亚 - 伽马（Pólya-gamma）增广的传统方法为何会导致贝叶斯推理难以处理。第 3.2 节阐述了我们基于充分统计量和偏微分方程（PDE）的替代建模方法的动机，并强调了模型的关键非可交换性和几何特性。最后，第 3.3 节提供了生成和推理的实现细节。

3.1 波利亚 - 伽马增广导致难以处理的推理

波利亚 - 伽马增广（Polson 等人，2013；Linderman 等人，2015）是一种强大的增广策略，通过引入额外的潜在变量来获得联合高斯似然，从而处理非高斯因素。Linderman 等人（2016）在 rSLDS 模型的原始公式中使用该方法，以允许通过消息传递进行可处理的推理。我们在附录 C.1 中表明，对于衰减函数 f 的适当选择，我们还可以利用波利亚 - 伽马增广来类似地处理由循环产生的非高斯因素。然而，由此产生的高斯增广对于每个时间步 n 呈线性增长，使其在计算上效率低下。幸运的是，如下所述，我们可以使用充分统计量和循环动态来规避这个问题。

3.2 用于高效参数化和非可交换性的 PDE 先验

为了弥补原始 dist-CRP 和 SLDS 模型简单组合所带来的挑战，我们依赖以下充分统计量来表示聚类分配：

在实际应用中，我们需要并行演化每个状态 j 的影响函数 w。对于任意固定的时间 t，我们通过使用平滑先验（见图 1C）来建立这些状态之间的关联，这是我们的建模选择。受（公式 6）中时间动态的启发，我们将充分统计量建模为根据偏微分方程（PDE）—— 热方程（见图 1）来演化。

我们强调“附近”一词的含义：在基于隐马尔可夫模型（HMM）的模型中，状态本身并没有固有的几何结构¹。相比之下，在这里，离散状态可以被解释为连续状态空间 J 上的点。我们可以将这个空间限制为离散数量的状态，并且我们将在下一小节中讨论如何在实践中做到这一点。

3.3 无限递归切换线性动态系统的实现与推理

接下来我们讨论上一节提出的偏微分方程先验在实现方面的若干问题。从现在起，我们假设数据在时间上是等间隔采样的，时间间隔为∆t，并省略时间变量 tₙ，仅考虑观测数据 yₙ。本文的代码基于 SSM 包（Linderman 等人，2020）开发，相关模块见附录 C.2。

基于有限差分法的偏微分方程先验实现
我们采用有限差分法求解满足热方程先验的充分统计量 w。定义 [wₙ₊₁]ⱼ₊₁ := w (j+∆j, (n+1)∆t) 为离散近似解。通过对时间导数应用前向差分、对空间二阶偏导数应用中心差分，可得到如下离散化方程：

其中，β = γ∆t/(∆j)²。为确保稳定性，我们通常要求∆t ≤ (∆j)²/(4γ)，并根据 γ（模型参数）和∆t（数据参数）相应调整∆j。系统可引入以下输入项驱动：(1) 期望的 κ₁zₙ=j 项，为系统提供自增强机制；(2) 由权重矩阵 R∈ℝᴷˣᴰ和偏置向量 r∈ℝᴷ构成的递归模块，编码历史内部状态信息。充分统计量 wₙ∈ℝᴷ的维度 K 是一个固定超参数。值得注意的是，热方程动态特性捕捉了距离相关 Chinese Restaurant Process (dist-CRP) 的特性，因此 K 实际上作为每条轨迹使用的简约状态数的上限。综上所述，wₙ的动态更新遵循：

其中扩散和衰减参数 γ ∝ β > 0，自增强参数 κ ∈ ℝ。公式（8）中向量 wₙ₊₁的第 j 个元素在经过必要的行归一化因子处理后，定义了转移矩阵 [Wₙ₊₁]ᵢⱼ的第 ij 个元素。动态模式 zₙ随后根据公式（5）进行采样，最终随机转移参数 α > 0。与 rSLDS 中的 K×K 矩阵相比，现在仅需几个标量参数 {α, γ, κ} 即可控制离散状态 zₙ的演化过程。公式（8）、（5）与公式（2）中的切换连续动态以及公式（3）中的发射模型共同定义了无限递归切换线性动态系统（irSLDS，见图 1F 中的图形模型）。关于模型在不同 K 值下的行为、参数 {α, γ, κ} 的影响以及这些参数的取值选择，详见附录 §D.3。

变分拉普拉斯期望最大化（Laplace-EM）推理方法

公式（8）中引入的充分统计量 wₙ由 zₙ₋₁和 xₙ₋₁确定性地决定，因此对后验推理没有额外贡献。我们采用 Zoltowski 等人（2020）提出的变分拉普拉斯期望最大化（Laplace-EM）方法进行模型推理。总体而言，我们对状态后验采用相同的结构化平均场近似 q (z₁ₜ) q (x₁ₜ)。对于连续状态后验，我们在最可能路径 x̂₁ₜ周围使用拉普拉斯近似（Paninski 等人，2009）。在该后验下给定连续状态轨迹后，可以计算转移矩阵 W₁ₜ来定义模型联合分布 p (y₁ₜ, z₁ₜ, x₁ₜ)。随后，通过在连续状态后验下局部最大化期望模型联合分布，得到离散状态后验近似。通常，在给定样本 x₁ₜ的条件下，离散状态 z 的后验因子图等价于非齐次隐马尔可夫模型（HMM）。此时可以使用常用工具，特别是维特比算法（Viterbi algorithm）来获取在 x̂₁ₜ条件下最可能的离散状态序列ẑ₁ₜ，我们将在下面进一步讨论。

离散状态几何结构

充分统计量 wₙ的支撑集定义为 {∆j, 2∆j, ..., K∆j}，它位于我们在 PDE 形式体系中使用的隐几何空间 J 上。在本文中，我们探索了一种利用这种几何结构的方法：定义中间状态。实际上，几何空间 J 的连续性允许在状态之间进行插值。对于离散状态 zₙ∈{1, ..., K} 上的任何分布，或其连续插值形式，可以定义各种感兴趣的统计量，如均值或众数。这些统计量可以在 J 中取连续值，而传统基于 HMM 的模型由于置换不变性无法定义此类统计量。特别地，给定最可能序列ẑ₁ₜ和我们的后验 q (zₙ₊₁ | zₙ, xₙ)，以及在 J 上的插值函数 q̃，我们将插值序列 z̃₁ₜ视为插值后的后验众数：

使用扫描运算并行化转移矩阵动态计算

在生成和推理过程中，无限递归切换线性动态系统（irSLDS）与传统递归切换线性动态系统（rSLDS）在计算复杂度上的核心差异在于：irSLDS 中离散状态 zₙ的转移矩阵 Wₙ自身具有递归动态特性（见公式 (8)，附录公式 (13)）。尽管 (r) SLDS 模型也具有时变转移矩阵，但它们不包含这种递归项。为了说明这种差异，考虑转移矩阵 Wₙ第 i 行的公式 (8)，rSLDS 中的等价表示为：

使用扫描运算并行化转移矩阵动态计算

在 rSLDS 中计算 zₙ的转移矩阵等价于对连续状态序列 x₁ₜ进行编码，这一过程可在不同处理器间并行分布。相比之下，irSLDS 中转移矩阵统计量 Wₙ的计算必须顺序进行，除了对 x₁ₜ的顺序编码外，还需额外执行 O (K³T) 次运算。幸运的是，注意到转移矩阵的动态方程（附录公式 13）是线性的。在这种情况下，根据 Blelloch（1990）的研究，转移矩阵的计算可以转化为扫描运算（有关这一概念的最新应用，请参见 Smith 等人（2023））。通过这种方法，我们可以高效地并行化转移矩阵的动态计算，并将计算复杂度降至 O (K³(T/L + log L))（在 L 个处理器上），从而达到与 rSLDS 模型相当的时间复杂度。有关并行扫描实现的代码和更多细节，请参见附录 C.2。

4 实验

我们通过使用 Zoltowski 等人（2020）提出的变分拉普拉斯期望最大化（Laplace-EM）方法最大化证据下界（ELBO）来训练模型。我们使用边际对数似然（LL）来比较模型性能。

4.1 合成 NASCAR 任务的验证

首先，我们发现 irSLDS 模型与真实生成模型 rSLDS 在生成动态轨迹的准确性上相当（见图 2D）。在轨迹的保留时间步（T=200）性能方面，两个模型的精度相近（见表 1）。当真实状态数 K=4 时，irSLDS 的平均性能更高；当状态数过完备（K=8）时，两者表现相当。而 SLDS 的性能不如预期 —— 该 NASCAR 任务最初用于测试 rSLDS 中递归项相对于标准 SLDS 的优势。尽管 SLDS 模型的测试 LL 相对较高（见表 1），但其难以准确刻画真实动态（见图 2D 左）。我们将这种生成过程的准确性称为 “动态精度”，通过计算学习动态与已知真实动态之间的均方误差（MSE）来量化，并对学习动态场的反射和旋转进行 MSE 最小化。结果显示 SLDS 的动态精度较差（见表 1），而 irSLDS 与真实 rSLDS 的精度相当。这验证了模型中递归项的有效性，同时表明 PDE 先验可模拟完全训练的 HMM 先验的性能和生成能力。

原始 NASCAR 任务具有周期性时间结构（一种不受置换不变性影响的属性）和固定状态数 K—— 两者均为平稳特性。在转向目标实验数据前，我们通过扩展 NASCAR 任务评估模型对非平稳挑战的处理能力。新任务中，动态在原始动态（图 2A）和一组新扩展动态（图 2B）之间交替，后者引入了两个额外的偏置动态，使离散状态真实数量在扩展动态块中从 K=4 变为 K=6。我们考虑 B∈{2,3} 个交替动态块，每个块的预期长度为 T/B。在该任务中，irSLDS 性能优于 rSLDS，rSLDS 又优于 SLDS（表 1）。分析 rSLDS 和 irSLDS 的学习解发现，irSLDS 的样本轨迹（图 2E）更接近数据的后验均值。此外，irSLDS 模型在动态切换后正确识别了两个新状态（图 2F 下），后验将这些状态归因于先前状态之间的过渡，使离散状态序列呈现轨迹穿越特性。相比之下，rSLDS 模型在切换到交替动态后错误地仅保留 4 个活跃状态（图 2F 上）。这表明 irSLDS 在非平稳任务中的优势，并举例说明了如何利用其离散状态几何结构。

4.2 irSLDS 揭示神经动态机制中的试次变异结构

接下来，我们将注意力转向国际脑实验室（IBL）的电生理数据集，该数据记录于小鼠执行感觉运动决策任务期间。在此任务中，小鼠通过向左或向右转轮来报告正弦光栅刺激的位置，任务难度由刺激对比度控制（图 3A）（Laboratory 等，2021）。我们使用 “Brainwide map” 数据发布中的 Neuropixels 探针记录（Laboratory，2022）。在分析时，我们将尖峰序列数据投影到主成分上以获取放电率，使其适用于高斯发射的状态空间模型分析（图 3B-C）（更多方法和数据细节见附录 §B，包括放电率和连续隐变量维度）。由于 irSLDS 是 rSLDS 的直接扩展，本节主要关注比较这两种模型。我们选择 K=8 个离散状态（K∈{2,4,8} 的测试边际 LL 值见附录表 3）。

首先，我们发现模型揭示了在任务相关事件时间点切换的离散隐状态。实际上，我们并未向模型提供刺激 onset（“Stim On”）或奖励等任务事件时间，也未提供运动起始时间（“First movement”）等行为测量数据。对于在时间步 t 发生的任务事件，我们在图 3D 中绘制了 irSLDS 估计的滞后 l 时 t+l 时间步的切换概率 p (switchₜ₊ₗ | eventₜ)，并与通过试次重排获得的基线概率 p (switchₜ) 进行比较。模型学习到的任务切换显著不同于随机水平，捕捉到了相关任务事件之前或之后的离散状态切换。因此，我们认为统计模型捕捉到了任务的相关动态。

其次，我们发现 irSLDS 揭示了试次间离散状态分布的差异。在图 3E 中，我们绘制了变分后验下最可能轨迹ẑ₁:ₜ中使用的活跃状态数量。我们对两种模型的数据拟合样条函数，并通过与常数函数比较的 F 检验最小 p 值确定（在不同平滑度和次数下计算）。尽管两种模型均显著不同于常数函数（p≪0.001），但 irSLDS 揭示的试次间离散状态数量曲线波动更大。rSLDS 模型的波动小得多，样条拟合效果更差（irSLDS 的 R² 为 0.46，而 rSLDS 为 0.16）。这表明在 irSLDS 模型下，试次间离散状态分布的波动更剧烈且系统，实验中段使用的离散状态数量显著多于开始或结束阶段。irSLDS 在保留测试试次上获得了略高的对数似然（见表 3），这表明这些波动很重要，而传统基于 HMM 的先验可能会忽略它们。

5 结论

在本研究中，我们通过在转移矩阵动态上引入输入驱动的热方程先验，对递归切换状态空间模型进行了扩展。这一改进形成了半马尔可夫离散状态过程，该过程既能捕捉每条轨迹中状态数量可变的分区随机过程行为，又能在离散状态上诱导连续几何结构。我们发现，尽管该过程的参数数量少于传统 HMM 模型，但其性能与原始模型相当甚至更优，同时为解析数据中的时变过程提供了新视角。我们先在合成任务上验证了模型有效性，随后将其应用于国际脑实验室（IBL）的电生理数据。IBL 拥有覆盖多次重复至全脑范围记录的海量数据集（Laboratory 等，2022；Laboratory，2022），而本文提出的建模方法为深入研究神经数据背后的时变过程奠定了基础。

更广泛的影响

本文提出的建模方法旨在增进我们对神经计算的理解。电生理数据分析可能对医学治疗和神经系统疾病的理解与治疗产生长远影响。然而，这些考量与本文呈现的初步分析和建模相距较远，我们预计这项工作不会立即产生社会影响。

局限性

我们强调几点局限性：首先，在实际实验中，第 3.3 节提出的用于并行计算的扫描操作实现不支持变分推理所需的自动微分，因此我们仅部分使用了扫描操作（具体方式详见附录 C.2）。其次，我们仅对有限数量的实验 session、探针位置和动物等电生理数据呈现了初步结果，并未试图对神经系统的计算机制做出普遍性论断，而是重点展示模型在揭示时间依赖过程中的特性。将研究扩展至上述数据模态是未来富有价值的工作方向。

计算资源

所有实验均在外部集群上运行。作为参考，对 IBL 电生理数据的单个 session 进行单模型训练（含部分多进程处理）可能耗时 12 小时，而合成 NASCAR 数据的单模型训练约需 30 分钟。

https://openreview.net/pdf?id=YIls9HEa52