使用 Attention Free Transformers 学习库普曼算子

# 使用无注意力前馈变换器学习Koopman算子 来源：https://arxiv.org/html/2606.23957 Mohammed Nagdi¹，Evangelos\-Marios Nikolados¹，Alexey Yermakov²，Mars Gao³，Nathan Kutz² 与 Filippo Menolascina¹ ¹英国爱丁堡大学生物工程研究所、工程学院及工程生物学中心 ²美国华盛顿大学西雅图分校电气与计算机工程与应用数学系 ³美国华盛顿大学西雅图分校电气与计算机工程与计算机科学与工程系 filippo\.menolascina@ed\.ac\.uk ###### 摘要 使用自动编码器学习Koopman算子能够在潜在空间中进行线性预测，但长时间跨度的展开常常偏离学习到的流形，导致在具有切换、连续谱或强瞬态的系统上出现相位和幅度误差。我们引入了两个互补组件，使Koopman预测器更加鲁棒。首先，我们添加了一个*无注意力潜在记忆*（AFT）模块，该模块聚合了过去一个短窗口的潜在状态，在每次Koopman更新之前产生一个修正的潜在状态。与多头注意力不同，AFT以线性时间运行，仅增加约≈30k个参数（\(3d^2 + T^2\)，少于匹配的多头注意力），同时捕获抑制误差发散所需的局部时间上下文。其次，我们提出了*动态重编码*：轻量级的在线变化点触发器（EWMA、CUSUM和顺序双样本检验），用于检测潜在漂移并将预测投影回自动编码器流形。在三个基准系统——Duffing振子、Repressilator、IRMA——上，我们的模型与Koopman自动编码器和匹配容量的多头注意力相比，一致地减少了误差累积。我们还与GRU和Transformer自动编码器进行了比较，从初始条件和50步上下文两种方式评估，发现Koopman+AFT（可选重编码）在保持更低推理延迟的同时，实现了显著更低的长时间跨度误差。我们报告了长达1000步的改进，以及关于触发器策略的消融实验。结果是一个快速、紧凑的预测器，能够在长时间跨度上保持在学习的流形上。

## 1 引言

Koopman算子通过将状态提升到一个演化是线性的可观测量空间，为使用线性工具分析非线性动力学提供了一种原则性方法（Koopman, 1931 (https://arxiv.org/html/2606.23957#bib.bib1)）。这种思想的神经实现——最常见的是Koopman自动编码器（KAE），它学习编码器\(\varphi\)、线性映射\(K\)和解码器\(\varphi^{-1}\)——通常能提供强大的单步精度，但在长时间展开中会出现漂移：振荡器中的相位滑动、幅度衰减或爆炸，以及轨迹偏离吸引子（Luschet al., 2018 (https://arxiv.org/html/2606.23957#bib.bib48)）。经验上，失败在以下情况下尤为突出：(i) 连续或混合谱（例如，无阻尼振荡器），(ii) 在亚稳态盆地之间切换，以及 (iii) 小误差会复合的瞬态区域。这促使我们设计机制：(a) 利用短期时间上下文来纠正局部误差，类似于Hankel DMD / HAVOK中的延迟嵌入思想（Arbabi and Mezic, 2017 (https://arxiv.org/html/2606.23957#bib.bib67); Bruntonet al., 2017 (https://arxiv.org/html/2606.23957#bib.bib68)），以及 (b) 在漂移变得灾难性之前周期性地将预测投影回学习到的流形。因此，长时间跨度的鲁棒性至关重要：保持在学习流形附近的预测器会累积更少的误差，并且更容易验证和用于下游控制。

**方法概述与直觉。** 我们在标准KAE的基础上增加两个部分：(i) *无注意力潜在记忆*（AFT）模块，聚合过去一个短窗口的潜在状态，在每次Koopman更新之前产生一个修正的潜在状态，在上下文长度上实现线性时间/内存开销，同时捕获驱动相位和幅度漂移的局部相关性（Zhaiet al., 2021 (https://arxiv.org/html/2606.23957#bib.bib59)）。(ii) *动态重编码* 使用轻量级流式触发器（EWMA、CUSUM、顺序双样本以及简单的阈值/窗口检验）检测潜在漂移，并应用编码-解码-编码（E–D–E）投影，将预测“拉回”到自动编码器流形上（Roberts, 2000 (https://arxiv.org/html/2606.23957#bib.bib55); Moustakides, 1986 (https://arxiv.org/html/2606.23957#bib.bib54); Ross and Adams, 2012 (https://arxiv.org/html/2606.23957#bib.bib56)）。直观地说，AFT解决了*如何*步进——在通过\(K\)传播之前减少局部误差——而重编码解决了*在哪里*步进——限制累积漂移。这两个机制是正交的：一个*阻止*增长，另一个*限制*它。

##### 与先前工作的关系。 我们的方法建立在数据驱动的Koopman学习基础上，从使用固定字典的EDMD到具有线性循环瓶颈的学习潜在嵌入（Liet al., 2017 (https://arxiv.org/html/2606.23957#bib.bib42); Otto and Rowley, 2019 (https://arxiv.org/html/2606.23957#bib.bib41); Luschet al., 2018 (https://arxiv.org/html/2606.23957#bib.bib48)）。短时间延迟上下文长期以来一直被用于稳定预测（Hankel DMD, HAVOK）（Arbabi and Mezic, 2017 (https://arxiv.org/html/2606.23957#bib.bib67); Bruntonet al., 2017 (https://arxiv.org/html/2606.23957#bib.bib68)），这启发了我们的轻量级潜在记忆。与最近混合模型中使用的Transformer式注意力（Luet al., 2024 (https://arxiv.org/html/2606.23957#bib.bib63); Wanget al., 2022 (https://arxiv.org/html/2606.23957#bib.bib65)）相比，我们的无注意力模块实现了线性成本，同时针对驱动相位/幅度漂移的局部相关性。正交地，投影/一致性思想（Nayaket al., 2025 (https://arxiv.org/html/2606.23957#bib.bib69); Frionet al., 2025 (https://arxiv.org/html/2606.23957#bib.bib70); Noacket al., 2015 (https://arxiv.org/html/2606.23957#bib.bib72); Dylewskyet al., 2019 (https://arxiv.org/html/2606.23957#bib.bib73); Guanet al., 2024 (https://arxiv.org/html/2606.23957#bib.bib71)）启发了我们的编码-解码-编码“拉回”机制。对于触发，我们采用了经典的流式漂移检测器（EWMA、CUSUM、顺序双样本）（Roberts, 2000 (https://arxiv.org/html/2606.23957#bib.bib55); Moustakides, 1986 (https://arxiv.org/html/2606.23957#bib.bib54); Ross and Adams, 2012 (https://arxiv.org/html/2606.23957#bib.bib56)）。关于输入和控制的更广泛脉络（KIC、Koopman MPC、安全/验证）在附录B (https://arxiv.org/html/2606.23957#A2) 中讨论，同时还包括生物学和流体力学中的特定应用。

**基准测试。** 我们针对三个代表性的系统，它们以互补的方式考验长时间跨度的稳定性：(i) *Duffing振子*，在无外力、无阻尼的情况下，呈现出闭合轨道，在更高能量下在势阱之间切换，并且通常具有连续或混合的Koopman谱，这对线性预测器构成压力（Otto and Rowley, 2019 (https://arxiv.org/html/2606.23957#bib.bib41); Liet al., 2017 (https://arxiv.org/html/2606.23957#bib.bib42); Pan and Duraisamy, 2020 (https://arxiv.org/html/2606.23957#bib.bib39); Alford-Lagoet al., 2022 (https://arxiv.org/html/2606.23957#bib.bib40); Köhneet al., 2025 (https://arxiv.org/html/2606.23957#bib.bib38)）；(ii) *Repressilator*，一个合成三基因负反馈振荡器，具有经典的极限环（Elowitz and Leibler, 2000 (https://arxiv.org/html/2606.23957#bib.bib33)），广泛用于评估识别和控制（Boddupalliet al., 2019 (https://arxiv.org/html/2606.23957#bib.bib34); Sootlaet al., 2018 (https://arxiv.org/html/2606.23957#bib.bib35); Balakrishnanet al., 2022 (https://arxiv.org/html/2606.23957#bib.bib36); Perez-Carrascoet al., 2018 (https://arxiv.org/html/2606.23957#bib.bib37)）；以及 (iii) *IRMA*（*In vivo Reverse-engineering and Modelling Assessment*，体内逆向工程与建模评估），一个五基因酵母回路，作为建模和控制的基准构建（Cantoneet al., 2009 (https://arxiv.org/html/2606.23957#bib.bib46); Marucciet al., 2009 (https://arxiv.org/html/2606.23957#bib.bib45); Menolascinaet al., 2014 (https://arxiv.org/html/2606.23957#bib.bib43); di Bernardoeet al., 2011 (https://arxiv.org/html/2606.23957#bib.bib44)），代表了多基因调控动力学，其中深度Koopman方法已显示出潜力（Hasnainet al., 2019 (https://arxiv.org/html/2606.23957#bib.bib47)）。这三个系统分别涵盖了混合谱和切换（Duffing）、具有相位敏感性的干净振荡行为（Repressilator）、以及具有交织反馈的高维调控动力学（IRMA）。

**经验总结。** 我们在三个主要基准（Duffing、Repressilator、IRMA）上进行了评估，报告了MSE和一个长时间跨度*平均累积绝对误差*（MCAE），该指标对误差累积敏感。潜在记忆模块在这些系统上优于匹配容量的MHA（4和10个头），并且将其与动态重编码结合，在切换和反馈丰富的系统（Duffing、IRMA）上产生了最鲁棒的展开，而纯AFT在干净的极限环（Repressilator）上仍然最优。GRU和Transformer自动编码器，从初始条件和50步上下文两种方式评估，在长时间跨度上表现不佳，尽管它们增加了上下文。

##### 贡献。
- • **用于Koopman预测的无注意力潜在记忆。** 一个线性时间、低开销的模块（Zhaiet al., 2021 (https://arxiv.org/html/2606.23957#bib.bib59)）聚合了潜在状态的历史短窗口，在每次Koopman更新之前产生一个修正的潜在状态，减少了长时间展开上的误差累积。
- • **通过流式变化检测实现动态重编码。** 一个编码-解码-编码投影，结合在线触发器（EWMA、CUSUM、顺序双样本、阈值/窗口），检测潜在漂移并将预测“拉回”到学习流形上（Roberts, 2000 (https://arxiv.org/html/2606.23957#bib.bib55); Moustakides, 1986 (https://arxiv.org/html/2606.23957#bib.bib54); Ross and Adams, 2012 (https://arxiv.org/html/2606.23957#bib.bib56)）。
- • **在三个代表性系统上的评估和消融实验。** 在Duffing（无外力、无阻尼）、Repressilator和IRMA上，潜在记忆优于匹配的MHA；潜在记忆 + 重编码在200/500/1000步范围内实现了最低的MSE；并且增益在Koopman算子大小变化时仍然保持。

参见图注

图1：具有AFT和动态重编码的Koopman自动编码器的工作流程。(a) 来自Duffing振子的采样轨迹作为输入。(b) 核心的**Koopman自动编码器**通过最小化重构、线性化和预测损失来学习线性潜在表示。(c) 预测过程使用带有AFT注意力的**动态重编码**模块来细化潜在状态（\(z_t \to \tilde{z}_t\)），然后由学习的Koopman算子 \(K\) 演化。(d) 最终输出显示匹配参考动力学的预测轨迹。

## 2 方法

### 2.1 基线Koopman自动编码器（KAE）

##### 模型。 令 \(x_t \in \mathbb{R}^p\) 表示时刻 \(t\) 的观测状态，并令 \(\varphi: \mathbb{R}^p \to \mathbb{R}^d\) 和 \(\varphi^{-1}: \mathbb{R}^d \to \mathbb{R}^p\) 为一对编码器/解码器，映射到 \(d\) 维潜在空间。KAE假设一个由学习的Koopman矩阵 \(K \in \mathbb{R}^{d \times d}\) 支配的*线性*潜在演化：
\[
z_t = \varphi(x_t), \qquad z_{t+1} = K z_t, \qquad \hat{x}_t = \varphi^{-1}(z_t),
\tag{1}
\]
因此，从初始潜在状态 \(z_0\) 开始的 \(i\) 步展开为 \(z_i = K^i z_0\)，解码后的预测为 \(\hat{x}_i = \varphi^{-1}(K^i \varphi(x_0))\)。我们使用标准的线性循环瓶颈架构（Otto and Rowley, 2019 (https://arxiv.org/html/2606.23957#bib.bib41); Luschet al., 2018 (https://arxiv.org/html/2606.23957#bib.bib48)）并端到端地学习 \((\varphi, \varphi^{-1}, K)\)。

##### 训练损失。 给定一个输入片段 \((x_0, \ldots, x_T)\)，我们最小化以下各项的加权和：(i) 重构误差，(ii) 潜在空间的线性一致性，(iii) 展开上的解码预测误差，以及 (iv) 对 \(K\) 的单式正则化以抑制谱的爆炸/消失（参见 Enyeart and Lin, 2024 (https://arxiv.org/html/2606.23957#bib.bib60)）：
\[
\begin{aligned}
\mathcal{L} &= \alpha_1 (\mathcal{L}_{\mathrm{recon}} + \mathcal{L}_{\mathrm{pred}}) + \mathcal{L}_{\mathrm{lin}} + \alpha_2 \mathcal{L}_{\mathrm{unitary}}, \tag{2a}\\
\mathcal{L}_{\mathrm{recon}} &= \frac{1}{T+1} \sum_{t=0}^T \|x_t - \varphi^{-1}(\varphi(x_t))\|_2^2, \tag{2b}\\
\mathcal{L}_{\mathrm{lin}} &= \frac{1}{T} \sum_{i=1}^T \|\varphi(x_i) - K^i \varphi(x_0)\|_2^2, \tag{2c}\\
\mathcal{L}_{\mathrm{pred}} &= \frac{1}{T} \sum_{i=1}^T \|x_i - \varphi^{-1}(K^i \varphi(x_0))\|_2^2, \tag{2d}\\
\mathcal{L}_{\mathrm{unitary}} &= \tfrac{1}{d^2} \|K K^\top - I\|_F. \tag{2e}
\end{aligned}
\]
权重 \(\alpha_1, \alpha_2 > 0\)。\(\mathcal{L}_{\mathrm{recon}}\) 强制保留信息的自动编码，\(\mathcal{L}_{\mathrm{lin}}\) 鼓励潜在轨迹与 \(K\) 的幂次的一致性，\(\mathcal{L}_{\mathrm{pred}}\) 衡量解码后的多步精度。单式惩罚轻微地将 \(K\) 偏向于接近正交，以提高长时间跨度的稳定性（Enyeart and Lin, 2024 (https://arxiv.org/html/2606.23957#bib.bib60)）。我们通过从 \(x_0\) 开始展开方程 (1) 共 \(T\) 步来训练，对同一片段计算所有四个损失。方程 (2) 中的公式与常见的KAE实践一致（Otto and Rowley, 2019 (https://arxiv.org/html/2606.23957#bib.bib41); Luschet al., 2018 (https://arxiv.org/html/2606.23957#bib.bib48)），同时明确了稳定性先验。

### 2.2 无注意力潜在记忆（AFT）

##### 设置。 为了缓解局部相位/幅度漂移，我们通过一个轻量级潜在记忆来增强KAE，该记忆在每次Koopman步骤之前聚合过去 \(T\) 个潜在状态（这里，\(T\) 表示*AFT上下文长度*，而不是方程 (2) 中的训练片段长度）。令时刻 \(t\) 的潜在历史为 \(H_t = [z_{t-T}, \ldots, z_{t-1}] \in \mathbb{R}^{T \times d}\)（我们使用因果索引且 \(T \ll\) 展开长度）。AFT模块将 \(H_t\) 映射到一个修正的潜在状态并更新：
\[
\tilde{z}_{t-1} = \mathrm{AFT}(H_t), \qquad z_t = K \tilde{z}_{t-1},
\tag{3}
\]
因此Koopman传播器推进的是一个*修正后*的潜在状态。

##### 计算。 这个变体是多头注意力（MHA）的即插即用替代品，可以被认为是一种逐元素的线性注意力机制。我们使用Zhai等人（2021 (https://arxiv.org/html/2606.23957#bib.bib59)）引入的AFT-full变体，其中给定Koopman子空间 \(Z_t\) 中 \(x_t\) 的潜在表示。我们应用学习的线性映射 \(W_Q, W_K, W_V \in \mathbb{R}^{d \times d}\)，
\[
Q_t = z_{t-1} W_Q \in \mathbb{R}^d, \quad K_t = H_t W_K / \sqrt{d_{\mathrm{model}}} \in \mathbb{R}^{T \times d},
\]