Training & Optimization

Production training pipeline using CatBoost and persistent Bayesian scaling

Data Science

What this does: Finds the best model settings automatically using Bayesian optimization, improving prediction accuracy without manual trial-and-error — saving days of data science effort per training cycle.

The system uses a production-grade CatBoost pipeline with persistent Bayesian optimization, designed for high-throughput scaling and reproducible experiments.

​

Training Pipeline

1. Data Preparation

   Features are engineered with time-based lags and rolling windows. Data is split using a chronological 80/20 split (preserving temporal order) to prevent look-ahead bias, ensuring the model is always evaluated on future data relative to the training set.

2. Hyperparameter Optimization

   Optuna explores the search space using a TPESampler. Unpromising trials are pruned early using a MedianPruner.

3. Production Model Selection

   The trial with the lowest RMSE is selected as the winner. The final model is retrained on the full dataset using the best hyperparameters found.

​

Persistent Optimization

Unlike standard in-memory optimization, our system persists the entire state of the hyperparameter search to a PostgreSQL database.

​

Key Benefits

1. Resume Capability: If the pipeline crashes, it resumes from where it left off.
2. Parallel Execution: Multiple Airflow workers can run trials concurrently.
3. Historical Analysis: Query the database later to analyze hyperparameter trends.

​

Implementation

airflow_dag_enhanced.py

# Unique study name ensures all runs contribute to the same optimization history
study = optuna.create_study(
    study_name="demand_forecasting_study_v2",
    storage="postgresql://airflow:airflow@postgres/optuna",
    load_if_exists=True,
    direction="minimize"
)

​

Configuration

Show Hyperparameter Search Space

Defined in app-ml/conf/base/parameters.yml. We optimize 5 key CatBoost parameters:

parameters.yml

catboost_search_space:
  learning_rate: { min: 0.01, max: 0.3 }
  depth: { min: 4, max: 12 }
  l2_leaf_reg: { min: 1.0, max: 10.0 }
  n_estimators: { min: 100, max: 500 }
  min_data_in_leaf: { min: 5, max: 50 }

Show Pruning Strategy

We use MedianPruner to stop unpromising trials early.

• Startup Trials: 5
• Warmup Steps: 10

Persistence: Optuna studies are saved to postgresql://airflow:airflow@postgres:5432/optuna.

​

Advanced Training Modes

The DAG supports two additional training strategies, enabled via parameters.yml:

Multi-Objective

Finds Pareto-optimal solutions balancing multiple metrics (RMSE and MAE).

• Trials: 50 (configurable via multi_objective_trials)
• Weights: RMSE 60%, MAE 40%
• Enable: advanced_features.enable_multi_objective: true

Ensemble

Combines multiple algorithms (CatBoost, XGBoost, LightGBM, RandomForest) using stacking.

• Method: Stacking (configurable via ensemble_type)
• Enable: advanced_features.enable_ensemble: true

​

Model Comparison

The pipeline also trains a Random Forest baseline via the demand_modeling Kedro pipeline and compares it head-to-head against CatBoost using RMSE, MAE, and MAPE. The winning model is selected automatically.

​

Evaluation Metrics

We prioritize RMSE for optimization, but monitor multiple metrics for business context.

RMSE

Root Mean Squared Error: The primary objective function.

$$RMSE = \sqrt{\frac{1}{n} \sum_{t=1}^{n} (y_t - \hat{y}_t)^2}$$

Why RMSE? It penalizes larger errors more heavily.

MAPE (Zero-Safe)

Mean Absolute Percentage Error:

To handle zero-demand hours, we replace actual zeros with 0.01 to prevent division-by-zero.

$$MAPE = \frac{100\%}{n} \sum_{t=1}^{n} \left| \frac{y_t - \hat{y}_t}{\text{where}(y_t=0, 0.01, y_t)} \right|$$

MAE

Mean Absolute Error: The average magnitude of errors.

$$MAE = \frac{1}{n} \sum_{t=1}^{n} |y_t - \hat{y}_t|$$

​

Optuna Visualization

Visualization

import optuna.visualization as vis

# Plot optimization history
vis.plot_optimization_history(study)

# Plot parameter importance
vis.plot_param_importances(study)

​

MLflow Integration

All training runs are automatically logged to the MLflow Tracking Server.

• Model Registry: Models are registered under demand_forecasting_catboost.
• Metrics: cv_rmse_mean, final_rmse, final_mae, final_mape.
• Artifacts: Code snapshots, regressor.pickle.

---

​

Related Pages

Parameters

Full search space configuration and feature engineering settings.

Experiment Tracking

MLflow integration for logging metrics, artifacts, and model registry.