Zero-Shot Probabilistic Forecasting with Lag-Llama

This guide shows you how to use the pre-trained Lag-Llama foundation model to generate probabilistic forecasts for a univariate time series without requiring explicit training (zero-shot). This is ideal for quick benchmarking or initial predictions on diverse datasets.

Prerequisites

• Chronax installed (pip install chronax).
• The target series y must be a 1-D jnp.ndarray of type float32.
• You must specify a context_length appropriate for your data, as this significantly impacts zero-shot performance.

Steps

1. Prepare the data and imports

Import necessary libraries and create a sample time series y. Lag-Llama requires the input data to be standardized as jnp.ndarray.

import jax.numpy as jnp
from chronax.models import LagLlama

# Example data (e.g., 100 observations)
y = jnp.sin(jnp.linspace(0, 10 * jnp.pi, 100)) + jnp.array(jnp.arange(100) * 0.1, dtype=jnp.float32)
y = y.astype(jnp.float32)

# Define the forecast horizon
H = 24

2. Instantiate the Lag-Llama model

Instantiate LagLlama. Since it is a foundation model, it automatically loads pre-trained weights for zero-shot use. You must specify the context_length, which determines how much historical data the model uses to condition the forecast.

# Start with a context length like 64 or 128 and tune this value
CONTEXT_LENGTH = 64 

model = LagLlama(context_length=CONTEXT_LENGTH)

3. Generate the zero-shot probabilistic forecast

Use the model.forecast method, passing the historical data y and the forecast horizon H. Since Lag-Llama is a probabilistic model, you should specify the level parameter to request prediction intervals.

# Request a 90% prediction interval (level=0.9)
forecast_output = model.forecast(
    y=y, 
    h=H, 
    level=0.9
)

4. Inspect the results

The forecast method returns a dictionary containing the point forecast ("mean") and the requested prediction intervals ("lo-90" and "hi-90").

mean_forecast = forecast_output["mean"]
lower_bound = forecast_output["lo-90"]
upper_bound = forecast_output["hi-90"]

print(f"Forecast length: {len(mean_forecast)}")
print(f"Mean forecast (first 5): {mean_forecast[:5]}")
print(f"Lower bound (first 5): {lower_bound[:5]}")

Full example

import jax.numpy as jnp
from chronax.models import LagLlama

# 1. Prepare the data
# Example data (100 observations)
y = jnp.sin(jnp.linspace(0, 10 * jnp.pi, 100)) + jnp.array(jnp.arange(100) * 0.1, dtype=jnp.float32)
y = y.astype(jnp.float32)
H = 24
CONTEXT_LENGTH = 64 

# 2. Instantiate the Lag-Llama model
model = LagLlama(context_length=CONTEXT_LENGTH)

# 3. Generate the zero-shot probabilistic forecast
# No explicit fit step is needed for zero-shot use.
forecast_output = model.forecast(
    y=y, 
    h=H, 
    level=0.9
)

# 4. Inspect the results
mean_forecast = forecast_output["mean"]
lower_bound = forecast_output["lo-90"]
upper_bound = forecast_output["hi-90"]

print(f"Forecast length: {len(mean_forecast)}")
print(f"Mean forecast (first 5): {mean_forecast[:5]}")
print(f"Upper bound (first 5): {upper_bound[:5]}")

Next steps

• Experiment with different context_length values in the LagLlama constructor to optimize zero-shot performance.
• Explore how to fine-tune LagLlama on your specific dataset using model.fit.
• Learn how to request multiple prediction intervals (e.g., 80% and 95%) simultaneously using the level parameter.
• Review the documentation for the model.forecast method.