Frequency-Aligned Knowledge Distillation for Lightweight Spatiotemporal Forecasting

Yuqi Li; Chuanguang Yang; Hansheng Zeng; Zeyu Dong; Zhulin An; Yongjun Xu; Yingli Tian; Hao Wu

Source-linked AI summary

Frequency-Aligned Knowledge Distillation for Lightweight Spatiotemporal Forecasting

Yuqi Li, Chuanguang Yang, Hansheng Zeng, Zeyu Dong, Zhulin An, Yongjun Xu, Yingli Tian, Hao Wu

arXiv:2507.02939v2Jun 27, 2025cs.LGcs.AIcs.CV

TL;DR

Spatiotemporal forecasters must model coupled local fluctuations and global evolution without the computational and memory burden of complex architectures. SDKD distills frequency-aligned multi-scale representations from a spectrally disentangled teacher into lightweight students, improving reported forecasting performance while reducing inference cost.

Problem
Complex hybrid spatiotemporal models are costly and generic distillation can ignore spectral bias, limiting efficient modeling of local variations and global trends.
Method
SDKD uses a convolution-Transformer teacher to disentangle high- and low-frequency latent features, then transfers aligned spectral representations to lightweight students.
Results
SDKD achieves or nearly matches teacher accuracy while reducing inference time across multiple datasets; on NS equation data, inference speed increases by 2.28× versus the teacher.
Takeaways & Limitations
Frequency-aligned distillation preserves multi-scale forecasting capability in lightweight students, including local details and global trends.

Abstract

from arXiv · show

Spatiotemporal forecasting tasks, such as traffic flow, combustion dynamics, and weather forecasting, often require complex models that suffer from low training efficiency and high memory consumption. This paper proposes a lightweight framework, Spectral Decoupled Knowledge Distillation (termed SDKD), which transfers the multi-scale spatiotemporal representations from a complex teacher model to a more efficient lightweight student network. The teacher model follows an encoder-latent evolution-decoder architecture, where its latent evolution module decouples high-frequency details and low-frequency trends using convolution and Transformer (global low-frequency modeler). However, the multi-layer convolution and deconvolution structures result in slow training and high memory usage. To address these issues, we propose a frequency-aligned knowledge distillation strategy, which extracts multi-scale spectral features from the teacher's latent space, including both high and low frequency components, to guide the lightweight student model in capturing both local fine-grained variations and global evolution patterns. Experimental results show that SDKD significantly improves performance, achieving reductions of up to 81.3% in MSE and in MAE 52.3% on the Navier-Stokes equation dataset. The framework effectively captures both high-frequency variations and long-term trends while reducing computational complexity. Our codes are available at https://github.com/itsnotacie/SDKD

1. Introduction

Spatiotemporal forecasting must capture both local high-frequency variations and global low-frequency evolution, but hybrid models and generic distillation can be costly and spectrally imbalanced. SDKD addresses this by transferring frequency-aligned multi-scale representations from a disentangled teacher to lightweight students.

Spatiotemporal forecasting couples high-frequency local patterns with low-frequency global evolution in tasks such as traffic and weather prediction.
Hybrid CNN-Transformer models capture local and long-range dependencies but incur prohibitive computational and memory costs.Their attention complexity scales as O(N^2d), with deep stacked structures further limiting deployment in resource-constrained environments.
Existing distillation methods can ignore spectral bias, suppressing low-frequency components and weakening generalizability across heterogeneous domains.
SDKD compresses complex teachers into lightweight students through frequency-aware representation alignment.The framework targets the efficiency-accuracy gap in spatiotemporal forecasting.
The teacher separates high-frequency details from low-frequency trends using convolution and Transformer modules to construct spectral priors.The convolution-Transformer serial design models local gradient responses and global attention patterns in the latent space.
Cross-scale spectral alignment guides students toward transient fluctuations and evolutionary trends, with reported gains for U-Net, ResNet, and MLP-Mixer models.On Weatherbench, U-Net with AEKD using AB Loss achieves MAE=0.8541, PSNR=31.4527, and SSIM=0.8812.

2. Related Work

Prior work combines spatial and temporal modeling, distills knowledge, or uses frequency-domain representations, but existing approaches often retain efficiency, cross-scale, or task-specific limitations. This paper positions spectral-aligned distillation as a distinct alternative.

Spatiotemporal Forecasting Models: ConvLSTM and PredRNN separate or sequence spatial-temporal modeling but struggle with cross-scale interactions between transients and trends.
Spatiotemporal Forecasting Models: Hybrid CNN-Transformer architectures improve local-global modeling but their quadratic sequence-length complexity limits long-horizon forecasting.
Knowledge Distillation Techniques: Regression distillation methods such as FitNet-R align features, but direct application to spatiotemporal forecasting can lose high-frequency information.
Frequency-Domain Representation Learning: Frequency-domain methods analyze spectral preferences, while Fourier operators model partial differential equations but have stated limitations.
Key Differences: The paper differentiates itself through a frequency-domain decoupled teacher and a spectral-aligned distillation mechanism for cross-scale spatiotemporal patterns.

3.1. Problem Formulation

The formulation treats forecasting as predicting future states from a spatiotemporal input sequence. It distills a costly hybrid teacher into a lightweight student using forecasting and spectral distillation objectives.

Given a spatiotemporal sequence X over H × W grids and C channels, the task predicts future states Y.
Traditional forecasting minimizes the discrepancy between the hybrid model output F(X) and the target Y.
The formulation distills F into a lightweight student G because Transformer quadratic complexity and deep CNNs limit deployment efficiency.
The objective combines forecasting loss with spectral distillation loss, using Ψ for spectral-feature extraction and λ to balance accuracy and efficiency.

3.2. Stage 1: Teacher Model Pretraining: Implicit Spectrum Decoupling

The teacher uses an encoder–latent evolution–decoder architecture that separates high-frequency details from low-frequency trends in latent space. CNN operations capture local variations, while self-attention models global trends before residual fusion and decoding.

Teacher architecture: The teacher maps inputs into latent space and processes them with parallel frequency-sensitive operations before reconstructing predictions through a decoder.The architecture uses convolution and Transformer components in latent evolution, followed by residual fusion and deconvolution.
High-frequency extraction: Convolution layers act as implicit high-frequency filters, using local receptive fields and gradient operations to capture abrupt changes and fine-grained details.The convolution kernel W_h has local receptive field size K, while differentiation order α and learnable coefficient a_α characterize the filtering behavior.
Low-frequency modeling: Self-attention acts as a low-pass filter that models low-frequency trends in meteorological evolution through global spatial interactions.The frequency variable ω and smoothing coefficient λ describe the filtering behavior.
Latent fusion: Residual fusion combines the high- and low-frequency latent components before the decoder reconstructs the forecast.The design explicitly preserves both spectral components for downstream prediction.
Physical interpretation: The teacher links high-frequency dynamics to viscous dissipation and low-frequency dynamics to convection, with experiments showing alignment to the Kolmogorov spectrum E(ω) ∝ω−5/3.This alignment is presented as evidence of physical interpretability in fluid-dynamics settings.

3.3. Stage 2: Lightweight Student Architecture

The student replaces the teacher’s heavy Transformer and deep convolutional structures with compact backbones while preserving multi-scale forecasting capacity and teacher-compatible input-output shapes.

Student design: The student can use a simple CNN, U-Net, or MLP backbone to reduce parameters and training memory while retaining capacity for multi-scale spatiotemporal patterns.The design goal is efficient deployment without removing the ability to represent local and global dynamics.
Backbone choices: ResNet, U-Net, and MLP variants provide progressively different efficiency and representation trade-offs for local features, cross-scale features, and low-resource settings.ResNet uses residual blocks, U-Net uses symmetric skip connections, and MLP reshapes the spatiotemporal grid for linear processing.
Efficiency: The lightweight backbones use O(Nd) or O(Kd^2) overhead instead of O(N^2d), with hidden dimension d set to 20%–30% of the teacher and models up to 5× smaller.The student preserves the teacher’s input-output shapes to support straightforward distillation.

3.4. Stage 3: Offline Training with Knowledge Distillation

Offline training transfers frequency-specific teacher representations to the student while jointly optimizing forecasting accuracy. In multi-teacher settings, A2D dynamically weights teacher gradients instead of averaging conflicting signals.

Offline distillation: The offline stage trains the student using one or more teachers pretrained with convolution-Transformer high-frequency versus low-frequency decoupling.The student receives supervision after teacher pretraining and lightweight architecture design.
Frequency-aligned distillation: The spectral transfer loss extracts high- and low-frequency components from the teacher’s latent representations to guide the student’s multi-scale forecasting.The method targets both transient fluctuations and evolutionary trends through frequency-specific supervision.
Training objective: The optimization jointly minimizes forecasting error and distillation loss, with λ balancing the two objectives.The frequency alignment term uses α to balance high- and low-frequency matching.
Multi-teacher distillation: For multiple teachers, A2D treats distillation as multi-objective optimization and solves teacher gradient weighting coefficients on a capped simplex during each mini-batch update.This replaces simple averaging when teachers specialize in different spatial, frequency, or physical domains.
Adaptive weighting: A2D decreases the influence of teachers whose gradient directions conflict strongly with others, preserving complementary information through adaptive gradient-level weighting.The resulting weights dynamically determine each teacher’s contribution to the student update.

3.5. Stage 4: Prediction with Lightweight Student

After training, the lightweight student generates predictions from a new input sequence in one forward pass. It removes the teacher’s quadratic-complexity Transformer blocks and deep convolution stacks while retaining multi-scale modeling through distillation.

3.5. Stage 4: Prediction with Lightweight Student: The deployed student G generates a prediction for each new input sequence Xnew in a single forward pass.The prediction is denoted Ŷnew = G(Xnew) ∈ R^H×W×C×∆.
3.5. Stage 4: Prediction with Lightweight Student: The lightweight architecture removes the teacher’s quadratic-complexity Transformer blocks and deep convolution stacks.
3.5. Stage 4: Prediction with Lightweight Student: 2.28× faster inference is achieved on the NS equation dataset compared with the teacher.The reported speed increase accompanies removal of the teacher’s Transformer blocks and deep convolution stacks.
3.5. Stage 4: Prediction with Lightweight Student: Frequency-aligned distillation preserves multi-scale modeling by combining local convolutions for high-frequency patterns with residual connections for global trends.The student does this without explicit frequency transforms and is intended to handle both steady-state and transient disturbances.

4. Experiment

Experiments evaluate SDKD across fluid, traffic, and weather forecasting benchmarks, comparing teacher–student models, distillation strategies, spectral behavior, and inference speed. The results show improved cross-scale prediction with substantial gains for selected methods and faster student inference on NS.

Experimental Settings: Four benchmarks cover RBC fluid dynamics, TaxiBJ+ traffic flow, WeatherBench meteorology, and Navier–Stokes forecasting.WeatherBench uses 12-hour inputs and forecasts across four meteorological variables; NS tests incompressible viscous-fluid dynamics.
Main Results: AEKD improves student performance across WeatherBench and TaxiBJ+, with U-Net achieving MAE 0.8541, PSNR 31.4527, and SSIM 0.8812 using AB Loss.ResNet and MLP-Mixer also improve in MAE and SSIM after AEKD distillation.
Spectral characteristics analysis: SDKD reduces both low- and high-frequency WeatherBench errors for U-Net, bringing the student closer to the teacher’s spectral performance.The reported mechanism aligns low-frequency Transformer attention and enhances high-frequency information using teacher latent features.
Distillation Method Analysis: CAMKD performs best on RBC, reaching MSE 0.0261, MAE 0.112, and SSIM 0.728 through multi-teacher distillation and spectral alignment.AVER-MKD reaches RBC MSE 0.0264, MAE 0.114, and SSIM 0.726, while the baseline reports MSE 0.0267 and MAE 0.114.
Inference Efficiency: On NS, distilled U-Net inference takes 0.0050 seconds versus 0.0115 seconds for SimVP, while RBC changes from 0.0192 to 0.0189 seconds.The speedup is larger on NS than on RBC.

5. Conclusion

The framework improves lightweight students’ multi-scale forecasting by aligning teacher spectral information across high-frequency details and low-frequency trends. It achieves or nearly matches teacher accuracy while reducing inference time across multiple datasets.

5. Conclusion: SDKD improves lightweight students’ ability to capture high-frequency details and low-frequency trends across multiple datasets.The framework explicitly aligns and transfers the teacher’s spectral information.
5. Conclusion: The framework achieves or nearly matches teacher accuracy while reducing inference time.The conclusion presents spectral decoupling distillation as applicable to spatiotemporal forecasting and multi-scale modeling.