Learning Rate Schedulers

Learning Rate Schedulers

Learning rate schedulers dynamically adjust the learning rate during the neural net training process, leading to improved convergence, faster training, and better model performance. Schedulers automate the process of adjusting the learning rate during training based on predefined rules or observed training performance. This reduces the need for painstaking trial-and-error experimentation.

You can optimize model convergence by using a higher learning rate early in training for larger steps in the parameter space, leading to faster initial progress. As training progresses and the model approaches convergence, you can use a lower learning rate to fine-tune the parameters and prevent overshooting the optimal solution, promoting stable convergence.

Using a fixed, small learning rate can trap the model in local minima or slow progress through flat regions of the loss landscape (saddle points). You can use a scheduler to introduce cycles of higher learning rates to escape these traps and explore the parameter space more effectively.

Conversely, large learning rates can lead to unstable training, where the loss function fluctuates wildly or even diverges. You can mitigate this by using a scheduler to gradually reduce the learning rate, which promotes a smooth and stable optimization process.

For validation and test sets, a well-chosen learning rate schedule enables you to fine-tune later stages of training. This helps the model generalize better to unseen data and leads to higher accuracy and better performance.

Supported schedulers

integrate.ai supports the following learning rate schedulers for VFL-SplitNN and VFL-GLM sessions:

Scheduler example

Add the learning rate scheduler (lrscheduler) parameters to the model_config for neural net training models (SplitNN or GLM).

model_config = {
    "strategy": {"name": "VflGlm", "params": {}},
    "model": {
        provider_name: {"params": {"input_size": len(provider_features), "output_activation": "exp"}},
        consumer_train_name: {"params": {"input_size": len(consumer_features), "output_activation": "exp"}},
    },
    "ml_task": {
        "type": "regression",
        "loss_function": "poisson",
        "params": {},
    },
    'lrscheduler': {"name": "onecyclelr", "params": {"max_lr": 0.1}},  # Scheduler
    "optimizer": {"name": "Adam", "params": {"learning_rate": 0.05}},
    "init_params": {consumer_train_name: baseline_init},
    "seed": 42,
    "calculate_training_metrics": True
}

model_config_params = {
    "lrscheduler.params.max_lr": [0.05, 0.1, 0.2, 0.3, 0.4],
}

Use the scheduler in batched sessions to quickly iterate through parameters.

Contents