STKAN：用于时空预测的Kolmogorov-Arnold Networks

# STKAN：基于柯尔莫哥洛夫-阿诺德网络的时空预测
来源：https://arxiv.org/html/2607.13108
，胡跃红 香港科技大学（广州）中国，钟思如 香港科技大学（广州）中国，乔思 香港科技大学（广州）中国，梁宇轩 香港科技大学（广州）中国，金广银 长安大学中国 中国地质大学中国

###### 摘要

现实交通数据呈现异质空间相关性和非线性时间动态，给准确的时空预测带来了重大挑战。现有方法已开发出日益复杂的图、注意力和分解架构，而底层非线性函数逼近器的影响却相对较少受到关注。本文提出STKAN，一种将泰勒多项式柯尔莫哥洛夫-阿诺德网络模块引入空间和时间令牌混合的时空预测架构。STKAN首先通过可学习的软节点组分配机制构建高层空间表示，应用组级空间混合，随后对压缩序列上的时间依赖性进行建模。进一步采用空间和时间自注意力层捕获长程交互。在五个交通预测基准上的实验表明，STKAN在测试设置中取得了有竞争力的性能，并优于评估的基于MLP的变体。这些结果表明，非线性函数逼近器的设计可以作为时空预测中架构设计的有用补充。

††copyright:none## 1. 引言

时空预测是AI4Transport生态系统中的基础支柱，驱动着智能交通技术（Chen等，2018 (https://arxiv.org/html/2607.13108#bib.bib59)；Wang等，2022 (https://arxiv.org/html/2607.13108#bib.bib60)；Lin等，2022 (https://arxiv.org/html/2607.13108#bib.bib61)）。通过从历史数据预测未来道路交通状况，AI智能体支持自动交通管理、路线优化和道路安全等服务（Wang等，2020 (https://arxiv.org/html/2607.13108#bib.bib62)；Deng等，2024 (https://arxiv.org/html/2607.13108#bib.bib63)；Han等，2022 (https://arxiv.org/html/2607.13108#bib.bib64)）。交通信号表现出空间异质性和复杂的时间动态。在空间上，区域间的交互复杂且非欧几里得。在时间上，信号表现出强周期性，同时夹杂着突变和非平稳模式。这些特性促使预测模型能够同时捕获空间依赖性和时间变化。

参见说明图1. 不同时空模型的示意图。

为应对这些挑战，研究人员开发了复杂的架构来显式建模这些依赖性。如图1 (https://arxiv.org/html/2607.13108#S1.F1)所示，基于图的方法通过假设预定义或学习到的图结构来执行预测。这些结构捕获节点间的空间相关性（Jin等，2023 (https://arxiv.org/html/2607.13108#bib.bib67)）。代表性工作如STGCN（Yu等，2018 (https://arxiv.org/html/2607.13108#bib.bib32)）和Graph WaveNet（Wu等，2019 (https://arxiv.org/html/2607.13108#bib.bib36)）利用图卷积和自适应邻接矩阵。同时，基于Transformer的方法利用注意力机制，建模时空令牌间的长程依赖（Liu等，2023 (https://arxiv.org/html/2607.13108#bib.bib56)）。此外，频谱分解技术如StemGNN（Cao等，2020 (https://arxiv.org/html/2607.13108#bib.bib53)）和STWave（Fang等，2023 (https://arxiv.org/html/2607.13108#bib.bib54)）专注于频域变换。像STID（Shao等，2022 (https://arxiv.org/html/2607.13108#bib.bib41)）这样的框架利用基于身份的嵌入。这些解耦方法旨在获得支持更准确预测的潜在时空表示。

尽管结构设计多样，许多时空架构仍使用带有固定激活函数的MLP作为默认的非线性映射模块。虽然MLP是富有表现力的通用逼近器，但其非线性通常预先选定并在单元间共享。这一观察引出一个补充性的研究问题：除了架构设计，非线性函数逼近器的选择在多大程度上影响时空预测性能？

本文研究非线性函数逼近器的选择是否是时空预测性能中一个未被充分探索的因素。我们提出STKAN，一种基于柯尔莫哥洛夫-阿诺德网络（KAN）的架构，它分解空间和时间依赖性，同时在空间和时间令牌混合映射中引入TaylorKAN层。STKAN提供了一个受控的架构环境，用于检验泰勒多项式KAN令牌混合器能否补充传统的时空建模组件。通道混合器和Transformer前馈网络仍基于MLP，因此该模型并未替换所有MLP组件。

我们的主要贡献总结如下：

- • 我们研究了非线性函数逼近器设计在时空预测中的作用，并引入了STKAN，它将泰勒多项式KAN映射融入空间和时间令牌混合模块中。
- • 我们开发了一个可学习的软节点组分配机制，以及专用的空间和时间混合模块。分组机制构建紧凑的空间表示，而两个令牌混合器分别沿空间组维度和时间维度建模交互。
- • 我们在五个交通预测基准上评估了STKAN。结果显示了有竞争力的预测精度，而现有的消融结果表明，在评估的配置中，TaylorKAN令牌混合器、自适应空间分组和注意力组件各自对整体模型有所贡献。

## 2. 相关工作

### 2.1. 时空预测

时空预测通过同时纳入时间动态和空间依赖性扩展了传统时间序列预测，例如在交通管理中，利用多个交通传感器的数据预测未来状况。早期深度学习方法结合卷积神经网络（CNN）和循环神经网络（RNN）来捕获空间和时间依赖性（Shi等，2015 (https://arxiv.org/html/2607.13108#bib.bib26)；Yao等，2018 (https://arxiv.org/html/2607.13108#bib.bib27)；Lai等，2018 (https://arxiv.org/html/2607.13108#bib.bib28)）。然而，基于网格的CNN可能无法有效处理非欧几里得空间关系，这促使了图卷积网络（GCN）（Defferrard等，2016 (https://arxiv.org/html/2607.13108#bib.bib29)；Kipf和Welling，2016 (https://arxiv.org/html/2607.13108#bib.bib30)）和时空图神经网络（STGNN）（Li等，2017 (https://arxiv.org/html/2607.13108#bib.bib31)；Yu等，2018 (https://arxiv.org/html/2607.13108#bib.bib32)；Li等，2023 (https://arxiv.org/html/2607.13108#bib.bib34)）的发展。这些模型，如DCRNN（Li等，2017 (https://arxiv.org/html/2607.13108#bib.bib31)）、ST-MetaNet（Pan等，2019 (https://arxiv.org/html/2607.13108#bib.bib33)）和DGCRN（Li等，2023 (https://arxiv.org/html/2607.13108#bib.bib34)），将GCN与RNN（Cho等，2014 (https://arxiv.org/html/2607.13108#bib.bib35)）集成，而其他模型如Graph WaveNet（Wu等，2019 (https://arxiv.org/html/2607.13108#bib.bib36)）和STGCN（Yu等，2018 (https://arxiv.org/html/2607.13108#bib.bib32)）则将GCN与门控时间卷积网络（TCN）结合。注意力机制也已被广泛采用于STGNN中（Zheng等，2020 (https://arxiv.org/html/2607.13108#bib.bib37)）。一些研究批评对预定义图的依赖，提出替代方案如AGCRN（Bai等，2020 (https://arxiv.org/html/2607.13108#bib.bib38)）和MTGNN（Wu等，2020 (https://arxiv.org/html/2607.13108#bib.bib39)），它们学习潜在图结构。最近的非GCN解决方案，如STNorm（Deng等，2021 (https://arxiv.org/html/2607.13108#bib.bib40)）和STID（Shao等，2022 (https://arxiv.org/html/2607.13108#bib.bib41)），进一步凸显了在预测任务中更好理解空间依赖性的需求。

### 2.2. 柯尔莫哥洛夫-阿诺德网络

KAN的动机来自于柯尔莫哥洛夫-阿诺德表示定理（Liu等，2024 (https://arxiv.org/html/2607.13108#bib.bib42)），该定理通过单变量函数的组合来表示多变量函数。基于这一范式，已经提出了几种变体，包括基于小波的WavKAN（Bozorgasl和Chen，2024 (https://arxiv.org/html/2607.13108#bib.bib44)）、基于泰勒多项式的TaylorKAN（Yu等，2025 (https://arxiv.org/html/2607.13108#bib.bib45)）、具有可训练雅可比基函数的fKAN（Aghaei，2025 (https://arxiv.org/html/2607.13108#bib.bib46)）以及采用径向基函数逼近B样条基的FastKAN（Li，2024 (https://arxiv.org/html/2607.13108#bib.bib47)）。

尽管KAN的动机来自柯尔莫哥洛夫-阿诺德表示定理，但该定理本身并不意味着有限KAN模型在经验预测任务中会普遍优于MLP。因此，在本文中，KAN被视作一种具有不同归纳偏置的替代非线性参数化方法，而非理论上保证的MLP替代品。

KAN最近在时间序列预测中引起了关注。先前的工作探索了基于KAN的模型用于时间序列分析（Vaca-Rubio等，2024 (https://arxiv.org/html/2607.13108#bib.bib65)）、混合专家KAN设计（Han等，2024 (https://arxiv.org/html/2607.13108#bib.bib43)）、面向符号回归的T-KAN和MT-KAN模型（Xu等，2024 (https://arxiv.org/html/2607.13108#bib.bib51)）、使用iTFKAN的协作时频学习（Liang等，2025 (https://arxiv.org/html/2607.13108#bib.bib48)）以及使用TimeKAN的频率分解学习（Huang等，2025 (https://arxiv.org/html/2607.13108#bib.bib49)）。BiLSTM-KAN将KAN层与双向LSTM骨干集成用于交通流预测（Chen等，2024 (https://arxiv.org/html/2607.13108#bib.bib66)）。STKAN与这些研究的不同之处在于，它使用泰勒多项式KAN层替换了空间和时间混合模块中的令牌混合映射，同时在通道混合和Transformer前馈组件中保留了传统的MLP。

## 3. 预备知识

时空预测是一个专门的多变量时间序列预测问题。给定历史输入序列 \(\mathbf{X}_{t-L_{\mathrm{in}}+1:t}=[\mathbf{X}_{t-L_{\mathrm{in}}+1},\ldots,\mathbf{X}_{t}]\in\mathbb{R}^{L_{\mathrm{in}}\times N\times C_{\mathrm{in}}}\)，目标是预测 \(\widehat{\mathbf{Y}}_{t+1:t+H}=[\widehat{\mathbf{Y}}_{t+1},\ldots,\widehat{\mathbf{Y}}_{t+H}]\in\mathbb{R}^{H\times N\times C_{\mathrm{out}}}\)。这里，\(L_{\mathrm{in}}\) 表示历史输入长度，\(H\) 表示预测视界，\(N\) 是空间节点数，\(C_{\mathrm{in}}\) 和 \(C_{\mathrm{out}}\) 是输入和输出通道维度。

## 4. 方法论

参见说明图2. 我们提出的STKAN框架概览。

在本文中，我们通过将TaylorKAN令牌混合器引入时空预测架构，提出STKAN来建模空间交互和时间动态。STKAN的整体架构如图2 (https://arxiv.org/html/2607.13108#S4.F2)所示。

### 4.1. 嵌入层

为捕获交通序列中的时空依赖性，我们使用全连接层嵌入原始输入 \(\mathbf{X}\in\mathbb{R}^{L_{\mathrm{in}}\times N\times C_{\mathrm{in}}}\)，得到 \(\mathbf{E}_{f}=\mathrm{FC}(\mathbf{X})\in\mathbb{R}^{L_{\mathrm{in}}\times N\times d_{f}}\)。我们使用一个可学习的空间嵌入 \(\mathbf{E}_{s}\in\mathbb{R}^{N\times d_{s}}\)。空间嵌入跨时间步共享，并在拼接前沿时间维度广播。为纳入时间周期性，我们使用两个嵌入字典：一个用于五分钟数据的日时间段嵌入，含 \(T_{\mathrm{tod}}=288\) 个槽；一个用于星期几嵌入，含 \(T_{\mathrm{dow}}=7\) 个类别。查表并广播到节点维度后，得到的嵌入分别记为 \(\mathbf{E}_{\mathrm{tod}}\) 和 \(\mathbf{E}_{\mathrm{dow}}\)。

### 4.2. 时间卷积块

所有嵌入沿特征维度拼接，形成初始表示：

\( (1) \mathbf{X}^{(0)} = \mathbf{E}_{f}\|\mathbf{E}_{s}\|\mathbf{E}_{\mathrm{tod}}\|\mathbf{E}_{\mathrm{dow}} \in \mathbb{R}^{L_{\mathrm{in}}\times N\times d_{h}}, \)

其中 \(d_{h}=d_{f}+d_{s}+d_{\mathrm{tod}}+d_{\mathrm{dow}}\)。为捕获局部时间上下文同时降低时间分辨率，我们定义一个补丁提取算子，沿时间轴应用步长为 \(s\) 的 \(1\times w\) 卷积：

\( (2) \mathbf{X}_{p} = \mathrm{PatchConv}(\mathbf{X}^{(0)}; w,s) \in \mathbb{R}^{L_{p}\times N\times d_{h}}, \)

其中 \(L_{p} = \lfloor (L_{\mathrm{in}}-w)/s \rfloor + 1\)。因此，\(\mathbf{X}_{p}\) 作为补丁级时空特征图，供后续混合器和注意力模块使用。

### 4.3. 时空KAN块

STKAN块结合了可学习空间分组、空间TaylorKAN令牌混合和时间TaylorKAN令牌混合。STKAN首先使用一个软分配矩阵将原始节点聚合成宏观空间令牌。聚合后的表示随后由空间KAN（SKAN）块处理。空间输出被映射回节点级特征，并传递给时间KAN（TKAN）块，后者对压缩后的时间维度上的依赖性进行建模。

**可学习空间分组**。为将原始空间节点转换为宏观空间令牌，我们引入一个未归一化的可学习参数矩阵 \(\mathbf{A}\in\mathbb{R}^{N\times G}\)。通过沿组维度应用softmax得到分配矩阵：

\( (3) \mathbf{G}_{n,g} = \frac{\exp(\mathbf{A}_{n,g})}{\sum_{g'=1}^{G}\exp(\mathbf{A}_{n,g'})}. \)

因此，对于每个节点 \(n\)，有 \(\sum_{g=1}^{G}\mathbf{G}_{n,g}=1\)。给定 \(\mathbf{X}_{p}\in\mathbb{R}^{L_{p}\times N\times d_{h}}\)，组级表示 \(\widetilde{\mathbf{X}}\in\mathbb{R}^{L_{p}\times G\times d_{h}}\) 计算如下：

\( (4) \widetilde{\mathbf{X}}_{\tau,g,c} = \sum_{n=1}^{N}\mathbf{G}_{n,g}\mathbf{X}_{p,\tau,n,c}. \)

**空间KAN块**。