FLAME: Flow Enhanced Legendre Memory Models for General Time Series Forecasting

This is the official repository of FLAME: Flow Enhanced Legendre Memory Models for General Time Series Forecasting. It has been accepted by NeurIPS 2026!

Introduction

FLAME is a family of extremely lightweight and highly capable time series foundation models. Based on the normalization-based forecasting head, it can support both the deterministic and probabilistic forecasting.

To our best knowldege, FLAME is the first time series foundation model possessing both lightweight backbones and generative prediction capabilities!

Architecture

FLAME adopts the Channel-Independent pretraining paradigm, and each variable is preprocessed through Instance Normalization to mitigate the value discrepancy. FLAME utilizes the Re-Norm to further mitigate the statistical differences between inputs and forecasts, and its backbone mainly consists of three modules: 1) Encoding, including Time Series Tokenization, Local-Perception, and MSA-Encoder, which tokenize the time series and enhance them through fusing the local environmental information with LegT; 2) Decoding, including LegS based SSD-Decoder and MCA-Enhancer, which utilize the SSD layers and MCA layers to make long-term inference ; 3) Flow-based Head, which leverages the Normalization Flow to support generative probabilistic forecasting, with both efficiency and accuracy.

Quickstart

We release all three versions of FLAME in different branches:

FLAME Small (2M) -- branch main & FLAME_Small
FLAME Base (6M) -- branch FLAME_Base
FLAME Large (10M) -- branch FLAME_Large

You need to install the following packages:

# pip install transformers[torch]

# pip install mamba-ssm[causal-conv1d]

# pip install zuko

To make deterministic or probabilistic forecasts, just follow:

from transformers import AutoModel, AutoConfig
import torch

model_path = "path/to/your/model"
config_path = "path/to/your/config"

config = AutoConfig.from_pretrained(config_path)

model = AutoModel.from_pretrained(model_path, config=config)
model.eval() 

 # The inputs need to be [batch_size, seq_len]. If multivariate, transform the inputs to [batch_size * n_vars, seq_len]
inputs = torch.randn(batch_size, seq_length)

# deterministic forecasting
with torch.no_grad():
    # output shape: [batch_size, 1, seq_len]
    outputs = model.generate(
        inputs=inputs,
        max_length=96,    
        revin=True,
        num_samples=1,
        inference_patch_len=48 # recommend to input the period length
    )


# probabilistic forecasting
with torch.no_grad():
    # output shape: [batch_size, 100, seq_len]
    outputs = model.generate(
        inputs=inputs,
        max_length=96,    
        revin=True,
        num_samples=100
    )