Learnable Koopman-Enhanced Transformer-Based Time Series Forecasting with Spectral Control

This paper introduces a unified family of learnable Koopman operator parameterizations that integrate linear dynamical systems theory with modern deep learning architectures to enable explicit spectral control, resulting in stable, interpretable, and high-performing time series forecasting models.

The Gist

Imagine you are trying to predict the future of a complex system, like the weather, the stock market, or how much electricity a city will need tomorrow. This is the job of Time Series Forecasting.

For a long time, scientists used simple rules (like "if it rained yesterday, it might rain today"). Then, we built powerful AI "black boxes" (like Transformers) that are great at spotting patterns but are hard to understand and sometimes act unpredictably.

This paper introduces a new tool called Learnable-DeepKoopFormer. Think of it as giving that powerful AI black box a GPS and a safety harness so it can navigate the future without crashing.

Here is the breakdown using simple analogies:

1. The Problem: The "Wild Horse" AI

Imagine a very smart horse (a modern AI model) that can run incredibly fast and see patterns in the forest. But, if you let it run without a bridle, it might get spooked, run in circles, or gallop off a cliff.

• The Issue: Current AI models are great at learning, but they can become unstable. If you ask them to predict 100 days into the future, they might start hallucinating wild numbers because their internal logic isn't "bouncy" enough to handle the long jump.

2. The Solution: The "Koopman" Safety Harness

The authors introduce a concept called the Koopman Operator.

• The Analogy: Imagine the weather is a chaotic dance. It's hard to predict the dancers' exact moves. But, if you look at the music (the underlying rhythm) instead of the dancers, the music follows simple, predictable rules.
• The Koopman operator is a mathematical trick that translates the chaotic "dance" of the data into a simple, predictable "music" (a linear system). It turns a messy, non-linear problem into a clean, straight line that is easy to control.

3. The Innovation: "Learnable" Control

Previous versions of this trick were like a rigid, pre-made harness. It was safe, but it couldn't adapt if the horse changed its personality.

• The New Idea: This paper creates Learnable harnesses. The AI gets to learn exactly how tight or loose the harness should be for each specific situation.
• They created four different types of "harnesses" (variants):
  1. Scalar-gated: A single master switch that controls the whole system's speed.
  2. Per-mode gated: Individual switches for every single part of the system (like adjusting the volume on every instrument in an orchestra separately).
  3. MLP-shaped: A smart, flexible neural network that shapes the rules dynamically.
  4. Low-rank: A simplified version that ignores the tiny, noisy details to focus on the big picture (like looking at a map from a high altitude rather than street level).

4. The "Spectral Control" (The Speed Limit)

The most important part of this paper is Spectral Control.

• The Analogy: Think of the AI's internal predictions as a ball rolling down a hill.
  • If the ball rolls too fast (unstable), it flies off the track.
  • If it stops too quickly (over-damped), it never reaches the destination.
  • Spectral Control is like putting a speed limit sign on the hill. It forces the AI to keep its internal "energy" within a safe, stable zone (between 0.3 and 0.8 on a scale of 0 to 1).
• This ensures the AI never explodes into chaos, but it also never freezes up. It keeps the "ball" rolling smoothly toward the future.

5. The Results: The "Goldilocks" Zone

The authors tested this on real-world data:

• Wind Speed: Predicting how hard the wind will blow.
• Air Pressure: Tracking weather systems.
• Cryptocurrency: Predicting wild crypto prices.
• Electricity: Guessing how much power a country needs.

What they found:

• Old AI (LSTMs, DLinear): Like the wild horse. Sometimes accurate, but often erratic and sensitive to small changes.
• Unconstrained New AI: Sometimes accurate, but prone to "crashing" (instability) when predicting far into the future.
• Learnable-DeepKoopFormer: This was the Goldilocks solution. It was:
  • Accurate: It predicted well.
  • Stable: It didn't crash or go crazy, even with difficult data like crypto.
  • Interpretable: Because of the "speed limit" (spectral control), scientists can actually look inside the AI and see why it made a prediction (e.g., "It's predicting a storm because the internal rhythm is slowing down").

Summary

This paper is about teaching powerful AI models to be responsible adults. Instead of just letting them learn whatever they want (which can be chaotic), the authors gave them a set of mathematical "training wheels" (the Koopman operator) that they can adjust themselves.

The result is a forecasting tool that is strong enough to handle complex, chaotic data (like the stock market) but stable enough to be trusted for critical tasks (like managing the power grid or predicting climate change). It's the difference between a race car with no brakes and a race car with a smart, adaptive braking system that keeps you fast but safe.

Technical Summary

1. Problem Statement

Time series forecasting is critical for applications in climate science, finance, and energy management. While modern deep learning architectures, particularly Transformers (e.g., PatchTST, Autoformer, Informer), have achieved state-of-the-art performance by modeling long-range dependencies, they face significant challenges:

• Lack of Dynamical Structure: Purely data-driven models often lack explicit physical or dynamical constraints, leading to poor robustness under distribution shifts.
• Instability: Unconstrained linear layers within deep networks can lead to unstable latent dynamics (explosive growth or vanishing signals) over long forecasting horizons.
• Interpretability: Standard deep models act as "black boxes," making it difficult to analyze the underlying temporal mechanisms or ensure stability guarantees.
• The Trade-off: Recent work (e.g., DLinear) suggests simple linear models can outperform complex Transformers, raising the question of how much architectural complexity is necessary while maintaining stability and interpretability.

The paper addresses the need for a framework that combines the expressiveness of Transformers with the theoretical stability and interpretability of linear dynamical systems, specifically leveraging Koopman operator theory.

2. Methodology: Learnable-DeepKoopFormer

The authors propose Learnable-DeepKoopFormer, a unified framework that integrates learnable Koopman operators into the latent space of Transformer backbones.

A. Core Architecture

The model follows an Encoder-Propagator-Decoder pipeline:

1. Encoder ($E_\theta$): A Transformer backbone (PatchTST, Autoformer, or Informer) processes an input window $X_t$ to produce a latent representation $z_t$.
2. Koopman Propagator ($K_\phi$): Instead of autoregressive rolling or complex recurrent steps, the latent state evolves linearly: $z_{t+1} = K_\phi z_t$. This operator is the core innovation.
3. Decoder ($D_\phi$): A linear mapping projects the propagated latent state back to the forecast horizon.

B. Learnable Koopman Variants

The key contribution is the parameterization of the linear propagator $K_\phi$ to allow learnable spectral control. The operator is factorized using an Orthogonal-Diagonal-Orthogonal (ODO) structure:
$$K_\phi = U_\phi \text{diag}(\Sigma_\phi) V_\phi^\top$$
where $U_\phi$ and $V_\phi$ are orthonormal matrices (learned via QR retraction), and $\Sigma_\phi$ contains the learnable spectral coefficients. The authors introduce four specific variants for generating these coefficients:

1. Scalar-Gated: A global affine gate $(\alpha, \beta)$ controls the entire spectrum, enabling tunable global damping.
2. Per-Mode Gated: Dimension-wise parameters $(\alpha_i, \beta_i)$ allow anisotropic temporal responses (different decay rates for different modes).
3. MLP-Shaped Spectral Mapping: A small neural network transforms raw spectral parameters before applying a squashing function, allowing flexible nonlinear reshaping of the spectrum.
4. Low-Rank Koopman: Restricts the operator to rank $r \ll d$, capturing dominant latent modes efficiently while reducing complexity.

Spectral Constraints:
To ensure stability, the spectral coefficients are passed through a smooth squashing map (logistic sigmoid) scaled by a bound $\rho_{max} < 1$:
$$\Sigma_{\phi, i} = \rho_{max} \cdot \sigma(\text{raw\_params})$$
This guarantees that the spectral radius $\rho(K_\phi) < 1$, ensuring exponential contraction and stability.

C. Training Objective

The model is trained end-to-end using a composite loss function:
$$L = L_{MSE} + \lambda_{Lyap} L_{Lyap}$$

• $L_{MSE}$: Standard Mean Squared Error for forecasting accuracy.
• $L_{Lyap}$: A Lyapunov-inspired penalty that penalizes latent energy growth ($\|z_{t+1}\|_P^2 - \|z_t\|_P^2$), further enforcing contractive dynamics and stability.

3. Key Contributions

1. Unified Parameterization: Introduces a family of four learnable Koopman operators that interpolate between strictly stable, constrained operators and unconstrained linear maps, providing explicit control over spectrum, stability, and rank.
2. Theoretical Guarantees: Proves that the proposed ODO parameterization ensures:
   • Spectral Stability: The spectral radius is strictly bounded below 1.
   • Invertibility: The operators are invertible (provided the spectrum is bounded away from zero), enabling well-posed backward propagation.
   • Lyapunov Consistency: The training objective biases the model toward operators satisfying discrete Lyapunov inequalities.
3. Comprehensive Benchmarking: Evaluates the framework across 21 model configurations (combining 3 backbones, 4 Koopman variants, and baselines) on five diverse real-world datasets (CMIP6 wind/pressure, ERA5 wind, Cryptocurrency, and Electricity generation).
4. Spectral Analysis: Provides the first systematic spectral analysis of Koopman-Transformers, visualizing eigenvalue trajectories, stability envelopes, and learned spectral distributions.

4. Experimental Results

The experiments compare Learnable-DeepKoopFormer against baselines including LSTM, DLinear, diagonal State-Space Models (SSM), and standard Transformers.

• Accuracy & Stability: Koopman-enhanced models consistently achieve lower and more tightly concentrated error distributions (MSE/MAE) compared to baselines. They are less sensitive to changes in patch length and forecast horizon.
• Spectral Behavior:
  • Constrained Variants: Show compact spectra in the range $\rho \approx 0.3–0.8$, ensuring stable, non-collapsed dynamics.
  • Learnable Variants: Occupy a similar range but with slightly broader distributions, successfully balancing expressiveness and stability.
  • Unconstrained Variants: Often collapse to very small radii ($\rho < 0.2$), leading to overly damped dynamics, or exhibit heavy tails indicating instability.
  • SSM Baseline: Shows occasional unstable modes (spectral radius $>1$) despite good finite-horizon fit.
• Domain Performance:
  • Climate/Weather: Outperforms baselines in wind speed and pressure forecasting, demonstrating robustness to non-stationarity.
  • Finance: Handles high volatility in cryptocurrency data better than recurrent models, with constrained variants showing the most stability.
  • Energy: Provides reliable long-horizon forecasts for electricity generation, where simple linear models (DLinear) and LSTMs showed overfitting or instability.

5. Significance

This work bridges the gap between operator-theoretic dynamics and modern deep learning.

• Interpretability: By explicitly controlling the spectrum, the latent dynamics become interpretable (e.g., identifying dominant modes and stability margins).
• Robustness: The spectral constraints prevent the "explosion" or "vanishing" of signals common in long-horizon forecasting, making the models suitable for high-stakes applications like climate risk assessment and energy planning.
• Efficiency: The low-rank variants demonstrate that high-dimensional dynamics can be captured efficiently without sacrificing stability.
• Future Impact: The paper establishes that spectral control is a powerful inductive bias for sequence models, suggesting a new direction for developing stable, interpretable, and theoretically grounded deep learning architectures.