Time series sales forecasting with IBM Granite

In this tutorial, we will explore time series forecasting using an IBM Granite® Time Series Foundation Model (TSFM) to predict retail sales. In addition to forecasting, we will cover key machine learning techniques such as few-shot fine-tuning, an efficient fine-tuning strategy by using only a portion of the designated training data. The sales data for this use case comes from the M5 datasets from the M-Competitions repository, provided as part of the long-running M-competition series for encouraging the research and development of forecasting methods. The aim of this tutorial is to forecast future sales aggregated by state, while showcasing how to use a pretrained TSFM for multivariate forecasting. In our analysis, we'll also use features available within the open source Granite Time Series Foundation Models toolkit, granite-tsfm.

In conjunction with thegranite-tsfm toolkit, this forecasting analysis leverages a TinyTimeMixer (TTM) model from the family of Granite TSFMs (available on HuggingFace and other open-source platforms). These TTMs are compact, pretrained models open sourced by IBM Research and capable of multivariate time series forecasting. The name TTM refers to the unique architecture of the models. With less than 1 million parameters, TTM introduced the notion of the first-ever "tiny" pretrained model for time series forecasting. TTM outperforms several popular benchmarks demanding billions of parameters in zero-shot and few-shot forecasting and can easily be fine-tuned for multivariate forecasts. Compared to traditional time series forecasting methods used in data science, Granite TSFMs have impressive zero-shot performance and support efficient fine-tuning for even more accurate predictions.

1.1: Install the TSFM library

The granite-tsfm library provides utilities for working with Time Series Foundation Models (TSFM). Here we retrieve and install the latest version of the Python library. The current Python versions supported are 3.9, 3.10, 3.11, 3.12.

1.2: Import packages

In addition to standard packages for data science, we'll use functionality from thetsfm_public/toolkit directory to prepare the data, fine-tune the model and generate forecasts.

To subset our data between training and test data, we'll use the functionsprepare_data_splits andselect_by_timestamp  . To visualize the data, we'll use theplot_predictions function. For preprocessing the data, theTimeSeriesPreprocessor class performs data transformations such as standardization and encoding categorical variables.

For fine-tuning, the toolkit uses a custom dataset type calledForecastDFDataset which is optimized for fast training and resulting forecasting by leveragingtorch . We'll also leverage theTrackingCallback class, thecount_parameters function and theoptimal_lr_finder function during fine-tuning.

To interact with the model, we'll use theTimeTimeMixerForPrediction class. Lastly, for forecasting we'll use theTimeSeriesForecastingPipeline class.

In addition to thetsfm_public toolkit, we'll use functionality fromtransformers andtorch for the fine-tuning step.

1.3: Specify configuration variables

Granite Time Series Models are sized with defined context lengths and forecast lengths. The context length will be sized accordingly with historical data, defining how many data points to look back into the past when making predictions. The forecast length specifies how many data points in the future to make predictions for. Here, we defineforecast_length andcontext_length to match the model we will select. It's important to note that these values must match the specification for the selected model.

Next, we declare the Granite Time Series Foundation Model, including the specific revision that we are targeting. In this time series analysis example, we will be working with daily data, so we choose a model suitable for that resolution—90 days of historical data to forecast the next 30 days. If a differentforecast_length orcontext_length makes sense for your data, the granite-timeseries TTM R2 card has several different revisions of the model available for various context lengths and prediction lengths.

As noted earlier, this notebook uses the M5 datasets from the official M-Competitions repository. You can read more about the competition and the dataset here.

2.1: Read in the data

Following initial data analysis, we observed that the original time series data includes hierarchy and product information. To prepare the data for this forecasting experiment, we will aggregate the sales by state into three separate time series. We'll use theprepare_data function in theM5_retail_data_prep.py file included here to download the datasets and prepare them as described.

Let's make sure we have access to theM5_retail_data_prep.py file (in an environment like colab, we need to download the file). Then, we simply run theprepare_data function to save the prepared dataset.

Following a typical data science workflow, we parse the resulting CSV file into apandas DataFrame . We also ensure that the timestamp column is a UTC datetime and drop two unnecessary columns of data points.

2.2: Set column identifiers

The next step for our time series analysis is to organize the columns in our sales data with the naming conventions required for input to our model.

In preparation for creating aTimeSeriesPreprocessor object, we'll set up thecolumn_specfiers dictionary here, declaring the names of thetimestamp_column , thetarget_column to be predicted, as well as indicatingcategorical_columns for encoding.

By specifying a list ofcols forcontrol_columns , we include the remaining columns in the forecasting dataset, allowing for the potential interactions of these exogenous variables with ourtarget_column , thesales values.

These column designations are important for the fine-tuning workflow, which permits multivariate forecasting by using the exogenous variables.

2.3: Train a preprocessor

Next, we set up a TimeSeriesPreprocessor, passing thecolumn_specifiers dictionary we created, as well as indicating that it should scale the data by usingsklearn's StandardScaler and ordinally encode thecategorical_columns we defined in the previous step.

Theselect_by_timestamp method is used to subset the training portion of the input data.

We train the scaling algorithm using thetsp.train method. Later, we'll use thepreprocess method on ourtsp preprocessor object to apply the scaling algorithm.

2.4: Split the data

For fine-tuning, we use the same data splits we defined earlier, but now we can include the extra columns as multivariate input.

We split the time series data into training, validation and test sets by using theprepare_data_splits function. The training and validation sets are used in the fine-tuning loop during model training, while the test set is used to evaluate the model performance after fine-tuning.

2.5: Create torch datasets

Next, we will construct three custom type ForecastDFDataset datasets by using the created data splitstrain_data ,valid_data andtest_data . We apply thepreprocess method from theTimeSeriesPreprocessor class to prepare the data before creating the dataset.

In comparison to apandas DataFrame , theForecastDFDataset type is specifically designed for fine-tuning and forecasting, leveragingtorch for faster performance. The HuggingFace Trainer API, which we'll be using for fine-tuning, can use these torch-based datasets.

2.6: Sample the data

Before beginning fine-tuning, we want to further subset the training and validation datasets, implementing a few-shot fine-tuning strategy. This strategy will be more efficient than if we fine-tuned the whole training dataset. Here, we sample the torch datasets produced from preceding steps, reducing the length of the dataset to 20% of the original size.

(2601, 520, 1398, 279)

Now we will focus on fine-tuning the pretrained model. The TinyTimeMixer architecture allows for fast fine-tuning, even on a CPU. In the following image from the IBM Research paper, the workflow for fine-tuning is illustrated in part (a).

In the diagram, the TTM architecture is depicted with 4 main components (a):

1. the TTM backbone composed of TSMixer blocks,
2. the TTM decoder, similarly architected, but 10–20% the size of the backbone,
3. the forecast head that produces forecasts and
4. an exogenous mixer, an optional component permitting multivariate forecasts.

The TTM decoder (2) and the forecast head (3) make up the TTM head. This head is typically retrained during fine-tuning, which is more efficient than fine-tuning the larger backbone component.

In contrast to the fine-tuning workflow, the pretrain workflow—also shown in part (a)—doesn't permit multivariate input, thus being unable to leverage exogenous variables in forecasting. The model is initially trained with univariate input, without considering interactions between variables, as indicated by the channel-independent methods in both the backbone and decoder components of the model for the pretrain workflow.

In the fine-tuning workflow, multivariate input is permitted and leveraged for its potential interactions with the target variable, by enabling channel mixing in the decoder and the optional exogenous mixer component.

You can read more about the steps for fine-tuning in the Granite docs.

3.1: Load the model

First, we use thefrom_pretrained method to load the TTM model (available on HuggingFace) by using the model and revision set from earlier. As indicated in our previously definedcolumn_specifiers dictionary, for this time series analysis, we have one target channel, several exogenous channels and one static categorical input. These exogenous channels, provided as additional columns beyond the target column in the fine-tuning workflow, might demonstrate interactions with the target channel that we can leverage for forecasting.

To provide these columns as input to the model, we provide theprediction_channel_indices ,exogenous_channel_indices andcategorical_vocab_size_list information to the model via theTimeSeriesPreprocessor objecttsp .

Note that we also enable channel mixing in the decoder by settingdecoder_mode="mix_channel" and forecast channel mixing by settingenable_forecast_channel_mixing=True leveraging the exogenous mixer component. This step allows the decoder to be tuned to capture interactions between the channels as well as to adjust the forecasts based on interactions with the exogenous, permitting multivariate forecasting.

After loading the model, we see a message reminding us to fine-tune the model before using it on a forecasting task.

3.2: Optional: Freeze the TTM backbone

Oftentimes, during fine-tuning we freeze the backbone component of the model, leaving these pretrained weights unchanged and we focus on tuning only the parameters in the decoder. This step reduces the overall number of parameters being tuned and maintains what the encoder learned during pretraining.

However, in this time series analysis, we found that performance was better when the backbone remained unfrozen—for other datasets one might prefer to freeze the backbone. We have disabled the backbone freezing code, but left it intact as an example of what might need to be done for other datasets.

3.3: Set hyperparameters

We'll use the Trainer API from HuggingFace for fine-tuning. To set up our fine-tuning, we need to specify values for hyperparameters. In the following code we set values fornum_epochs ,batch_size andlearning_rate which we'll soon pass asTrainingArguments when we create aTrainer object. We use the optimal_lr_finder function from thetsfm_public toolkit to find the optimallearning_rate for this dataset. Note that theseTrainingArguments are specific to this particular dataset.

OPTIMAL SUGGESTED LEARNING RATE = 0.000298364724028334

3.4: Train the model

Here we train the model on the historical training data by using the hyperparameters that we previously set. We create theTrainer object by using the specifiedTrainingArguments , anEarlyStoppingCallback criteria, an AdamW optimizer andOneCycleLR scheduler. We've set thenum_epochs to50 , but theEarlyStoppingCallback implemented here will stop the training after 10 epochs with no improvement. After theTrainer object is configured, we call thetrain method to perform fine-tuning and the evaluate method to provide metrics about the fine-tuning results.

Using learning rate = 0.000298364724028334

{'eval_loss': 0.3611494302749634,
'eval_runtime': 15.3951,
'eval_samples_per_second': 90.613,
'eval_steps_per_second': 0.715,
'epoch': 37.0}

4.1: Generate forecasts

We'll leverage the TimeSeriesForecastingPipeline fromgranite-tsfm to create forecasts by using our fine-tuned model. The following is a preview of these forecasts.

4.2: Evaluate the model

To assess whether the model produces accurate predictions for our time series analysis, we’ll evaluate the fine-tuned model on the originaltest_data (apandas dataset). Then, we'll quantify the forecast errors by defining acustom_metric to calculate the MSE (Mean Squared Error), RMSE (Root Mean Squared Error) and MAE (Mean Absolute Error).

Then, we'll calculate error metrics using the predefinedcustom_metric . The RMSE (Root Mean Squared Error) and MAE (Mean Absolute Error) metrics demonstrate low error for this forecasting problem on this dataset. While the sales values that we are predicting are in the tens of thousands, the RMSE and MAE metrics give errors hovering around USD 1,000. The MSE (Mean Squared Error) is higher, but that might be due to outlier predictions with high errors.

4.3: Plot the predictions vs. actuals

Finally, we use the plot_predictions function fromgranite-tsfm to plot the predictions or future values for sales forecasting vs. the actual values for some random samples of time intervals in test dataset, visualizing the forecast error. Using random samples of time intervals allows us to test how the model handled fluctuations in the data. We can also observe the past values in these plots that served as context length.

In this tutorial, we performed time series forecasting for a sales data use case. Our time series analysis demonstrates forecasting methods by using a Granite TSFM model. This model is a compact pretrained foundation model that has been shown to outperform other machine learning algorithms including ARIMA (AutoRegressive Integrated Moving Average) and LSTM (Long Short-Term Memory networks). While our use case focuses on sales forecasting, the Granite TSFM models can aid data scientists by forecasting values for numerous variables while being robust to outliers and seasonal variations in the data. Some of examples of these variables include weather, stock prices, flu cases, other temporal data and more. Time series predictions like these enable real world data-driven decision making.