xbe.at

Deep learning / Backpropagation and optimization

Exploring Learning Rate in Deep Learning with Python

10 min read, 15 sections

Contents
  1. What is Learning Rate?
  2. Impact of Learning Rate
  3. Types of Learning Rate Schedules
  4. Learning Rate Warmup
  5. Adaptive Learning Rate Methods
  6. Learning Rate and Batch Size Relationship
  7. Learning Rate Finder
  8. Cyclical Learning Rates
  9. One Cycle Policy
  10. Learning Rate and Model Convergence
  11. Learning Rate and Generalization
  12. Real-life Example: Image Classification
  13. Real-life Example: Natural Language Processing
  14. Learning Rate in Reinforcement Learning
  15. Additional Resources

What is Learning Rate?

Learning rate is a crucial hyperparameter in deep learning that determines the step size at each iteration while moving toward a minimum of the loss function. It controls how quickly or slowly a neural network model learns from the data.

import numpy as np
import matplotlib.pyplot as plt

def gradient_descent(learning_rate, iterations):
    x = np.linspace(-10, 10, 100)
    y = x**2  # Simple quadratic function
    
    current_x = 8  # Starting point
    trajectory = [current_x]
    
    for _ in range(iterations):
        gradient = 2 * current_x  # Derivative of x^2
        current_x = current_x - learning_rate * gradient
        trajectory.append(current_x)
    
    plt.plot(x, y)
    plt.plot(trajectory, [t**2 for t in trajectory], 'ro-')
    plt.title(f"Gradient Descent (LR = {learning_rate})")
    plt.show()

gradient_descent(0.1, 10)

Impact of Learning Rate

The learning rate significantly affects the training process. A too-high learning rate can cause the model to overshoot the optimal solution, while a too-low learning rate can result in slow convergence or getting stuck in local minima.

import torch
import torch.nn as nn
import torch.optim as optim

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(1, 1)
    
    def forward(self, x):
        return self.linear(x)

def train_model(learning_rate):
    model = SimpleModel()
    criterion = nn.MSELoss()
    optimizer = optim.SGD(model.parameters(), lr=learning_rate)
    
    X = torch.linspace(-10, 10, 100).reshape(-1, 1)
    y = 2 * X + 1 + torch.randn(X.shape) * 0.1
    
    losses = []
    for epoch in range(100):
        optimizer.zero_grad()
        outputs = model(X)
        loss = criterion(outputs, y)
        loss.backward()
        optimizer.step()
        losses.append(loss.item())
    
    plt.plot(losses)
    plt.title(f"Training Loss (LR = {learning_rate})")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.show()

train_model(0.01)
train_model(0.1)

Types of Learning Rate Schedules

Different learning rate schedules can be employed to improve training:

  1. Constant: Fixed learning rate throughout training
  2. Step decay: Reduce learning rate at predetermined intervals
  3. Exponential decay: Continuously decrease learning rate exponentially
  4. Cosine annealing: Cyclical learning rate following a cosine function
import numpy as np
import matplotlib.pyplot as plt

def constant_lr(initial_lr, epochs):
    return [initial_lr] * epochs

def step_decay(initial_lr, epochs, drop_rate=0.5, epochs_drop=10):
    return [initial_lr * (drop_rate ** (epoch // epochs_drop)) for epoch in range(epochs)]

def exp_decay(initial_lr, epochs, k=0.1):
    return [initial_lr * np.exp(-k * epoch) for epoch in range(epochs)]

def cosine_annealing(initial_lr, epochs, T_max=50):
    return [initial_lr * (1 + np.cos(np.pi * epoch / T_max)) / 2 for epoch in range(epochs)]

epochs = 100
initial_lr = 0.1

plt.figure(figsize=(12, 8))
plt.plot(constant_lr(initial_lr, epochs), label='Constant')
plt.plot(step_decay(initial_lr, epochs), label='Step Decay')
plt.plot(exp_decay(initial_lr, epochs), label='Exponential Decay')
plt.plot(cosine_annealing(initial_lr, epochs), label='Cosine Annealing')
plt.xlabel('Epochs')
plt.ylabel('Learning Rate')
plt.title('Learning Rate Schedules')
plt.legend()
plt.show()

Learning Rate Warmup

Learning rate warmup is a technique where the learning rate starts from a small value and gradually increases to the initial learning rate. This can help stabilize training in the early stages.

import numpy as np
import matplotlib.pyplot as plt

def warmup_lr(initial_lr, warmup_epochs, total_epochs):
    lr_schedule = []
    for epoch in range(total_epochs):
        if epoch < warmup_epochs:
            lr = initial_lr * (epoch + 1) / warmup_epochs
        else:
            lr = initial_lr
        lr_schedule.append(lr)
    return lr_schedule

initial_lr = 0.1
warmup_epochs = 10
total_epochs = 100

lr_schedule = warmup_lr(initial_lr, warmup_epochs, total_epochs)

plt.plot(lr_schedule)
plt.xlabel('Epochs')
plt.ylabel('Learning Rate')
plt.title('Learning Rate Warmup')
plt.show()

Adaptive Learning Rate Methods

Adaptive learning rate methods automatically adjust the learning rate during training. Popular algorithms include:

  1. AdaGrad
  2. RMSprop
  3. Adam (Adaptive Moment Estimation)
import torch
import torch.nn as nn
import torch.optim as optim

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(1, 1)
    
    def forward(self, x):
        return self.linear(x)

def compare_optimizers():
    X = torch.linspace(-10, 10, 100).reshape(-1, 1)
    y = 2 * X + 1 + torch.randn(X.shape) * 0.1
    
    optimizers = {
        'SGD': optim.SGD,
        'AdaGrad': optim.Adagrad,
        'RMSprop': optim.RMSprop,
        'Adam': optim.Adam
    }
    
    plt.figure(figsize=(12, 8))
    
    for name, opt_class in optimizers.items():
        model = SimpleModel()
        optimizer = opt_class(model.parameters(), lr=0.01)
        criterion = nn.MSELoss()
        
        losses = []
        for _ in range(100):
            optimizer.zero_grad()
            outputs = model(X)
            loss = criterion(outputs, y)
            loss.backward()
            optimizer.step()
            losses.append(loss.item())
        
        plt.plot(losses, label=name)
    
    plt.xlabel('Iterations')
    plt.ylabel('Loss')
    plt.title('Comparison of Optimization Algorithms')
    plt.legend()
    plt.show()

compare_optimizers()

Learning Rate and Batch Size Relationship

The learning rate and batch size are interconnected. Increasing the batch size allows for a larger learning rate, which can lead to faster convergence. This relationship is known as the "linear scaling rule."

import numpy as np
import matplotlib.pyplot as plt

def train_with_batch_size(batch_size, learning_rate, epochs=100):
    losses = []
    for _ in range(epochs):
        batch_loss = np.random.normal(loc=0.5, scale=0.1) / np.sqrt(batch_size)
        losses.append(batch_loss)
    return np.cumsum(losses) * learning_rate

batch_sizes = [32, 64, 128, 256]
learning_rates = [0.001, 0.002, 0.004, 0.008]

plt.figure(figsize=(12, 8))
for bs, lr in zip(batch_sizes, learning_rates):
    losses = train_with_batch_size(bs, lr)
    plt.plot(losses, label=f'Batch Size: {bs}, LR: {lr}')

plt.xlabel('Iterations')
plt.ylabel('Cumulative Loss')
plt.title('Learning Rate and Batch Size Relationship')
plt.legend()
plt.show()

Learning Rate Finder

The learning rate finder is a technique to determine a good initial learning rate. It involves training the model for a few iterations with exponentially increasing learning rates and plotting the loss.

import numpy as np
import matplotlib.pyplot as plt

def simulate_training(min_lr, max_lr, num_iterations):
    lrs = np.logspace(np.log10(min_lr), np.log10(max_lr), num_iterations)
    losses = []
    for lr in lrs:
        loss = np.random.normal(loc=1 - np.exp(-lr * 10), scale=0.1)
        losses.append(loss)
    return lrs, losses

min_lr, max_lr = 1e-5, 1
num_iterations = 100

lrs, losses = simulate_training(min_lr, max_lr, num_iterations)

plt.figure(figsize=(12, 6))
plt.semilogx(lrs, losses)
plt.xlabel('Learning Rate')
plt.ylabel('Loss')
plt.title('Learning Rate Finder')
plt.grid(True)
plt.show()

optimal_lr = lrs[np.argmin(losses)]
print(f"Suggested Learning Rate: {optimal_lr:.2e}")

Cyclical Learning Rates

Cyclical Learning Rates (CLR) involve cycling the learning rate between a lower and upper bound. This can help overcome saddle points and local minima.

import numpy as np
import matplotlib.pyplot as plt

def triangular_clr(initial_lr, max_lr, step_size, iterations):
    cycle = np.floor(1 + iterations / (2 * step_size))
    x = np.abs(iterations / step_size - 2 * cycle + 1)
    lr = initial_lr + (max_lr - initial_lr) * np.maximum(0, (1 - x))
    return lr

iterations = 1000
initial_lr = 0.001
max_lr = 0.1
step_size = 200

lr_schedule = [triangular_clr(initial_lr, max_lr, step_size, i) for i in range(iterations)]

plt.figure(figsize=(12, 6))
plt.plot(lr_schedule)
plt.xlabel('Iterations')
plt.ylabel('Learning Rate')
plt.title('Triangular Cyclical Learning Rate')
plt.show()

One Cycle Policy

The One Cycle Policy is a learning rate schedule that involves one cycle of increasing and then decreasing learning rate, combined with momentum oscillations.

import numpy as np
import matplotlib.pyplot as plt

def one_cycle_policy(max_lr, total_steps, pct_start=0.3):
    steps = np.arange(total_steps)
    step_size = int(total_steps * pct_start)
    
    lr_schedule = np.where(steps < step_size,
                           max_lr * (steps / step_size),
                           max_lr * (1 - (steps - step_size) / (total_steps - step_size)))
    
    momentum_schedule = np.where(steps < step_size,
                                 0.95 - 0.45 * (steps / step_size),
                                 0.95 - 0.45 * (1 - (steps - step_size) / (total_steps - step_size)))
    
    return lr_schedule, momentum_schedule

total_steps = 1000
max_lr = 0.1

lr_schedule, momentum_schedule = one_cycle_policy(max_lr, total_steps)

fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 10))
ax1.plot(lr_schedule)
ax1.set_ylabel('Learning Rate')
ax1.set_title('One Cycle Policy')

ax2.plot(momentum_schedule)
ax2.set_xlabel('Steps')
ax2.set_ylabel('Momentum')

plt.tight_layout()
plt.show()

Learning Rate and Model Convergence

The learning rate significantly impacts model convergence. A well-chosen learning rate leads to faster and more stable convergence, while a poor choice can result in slow learning or divergence.

import numpy as np
import matplotlib.pyplot as plt

def simulate_convergence(learning_rates, iterations=1000):
    losses = {}
    for lr in learning_rates:
        loss = 100
        loss_history = []
        for _ in range(iterations):
            gradient = np.random.normal(0, 1)
            loss -= lr * gradient
            loss = max(0, loss)  # Ensure non-negative loss
            loss_history.append(loss)
        losses[lr] = loss_history
    return losses

learning_rates = [0.001, 0.01, 0.1, 1.0]
convergence_data = simulate_convergence(learning_rates)

plt.figure(figsize=(12, 6))
for lr, loss_history in convergence_data.items():
    plt.plot(loss_history, label=f'LR = {lr}')

plt.xlabel('Iterations')
plt.ylabel('Loss')
plt.title('Impact of Learning Rate on Convergence')
plt.legend()
plt.yscale('log')
plt.show()

Learning Rate and Generalization

The learning rate can affect the model's ability to generalize. A very high learning rate might lead to overfitting, while a very low learning rate might result in underfitting.

import numpy as np
import matplotlib.pyplot as plt

def generate_data(n_samples=100):
    X = np.linspace(0, 10, n_samples)
    y = 2 * X + 1 + np.random.normal(0, 2, n_samples)
    return X, y

def train_model(X, y, learning_rate, epochs=1000):
    w, b = 0, 0
    for _ in range(epochs):
        y_pred = w * X + b
        error = y - y_pred
        w += learning_rate * np.mean(error * X)
        b += learning_rate * np.mean(error)
    return w, b

X, y = generate_data()
X_test = np.linspace(0, 10, 100)

plt.figure(figsize=(15, 5))
learning_rates = [0.0001, 0.01, 0.5]

for i, lr in enumerate(learning_rates, 1):
    w, b = train_model(X, y, lr)
    y_pred = w * X_test + b
    
    plt.subplot(1, 3, i)
    plt.scatter(X, y, alpha=0.5)
    plt.plot(X_test, y_pred, 'r-', label='Prediction')
    plt.title(f'LR = {lr}')
    plt.legend()

plt.tight_layout()
plt.show()

Real-life Example: Image Classification

In image classification tasks, the learning rate significantly impacts the training of deep neural networks. Let's consider a simple CNN for MNIST digit classification.

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.relu(self.conv2(x))
        x = torch.flatten(x, 1)
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

def train(model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = nn.functional.cross_entropy(output, target)
        loss.backward()
        optimizer.step()

# Usage example (not executed)
# model = SimpleCNN().to(device)
# optimizer = optim.Adam(model.parameters(), lr=0.001)
# train(model, device, train_loader, optimizer, epoch)

Real-life Example: Natural Language Processing

In NLP tasks, learning rate tuning is crucial for training large language models. Here's a simplified example of a recurrent neural network for sentiment analysis.

import torch
import torch.nn as nn

class SentimentRNN(nn.Module):
    def __init__(self, vocab_size, embed_size, hidden_size, output_size):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.rnn = nn.LSTM(embed_size, hidden_size, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = self.embedding(x)
        _, (hidden, _) = self.rnn(x)
        out = self.fc(hidden.squeeze(0))
        return torch.sigmoid(out)

def train_step(model, optimizer, criterion, inputs, labels):
    model.train()
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()
    return loss.item()

# Usage example (not executed)
# model = SentimentRNN(vocab_size, embed_size, hidden_size, output_size)
# optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# criterion = nn.BCELoss()
# loss = train_step(model, optimizer, criterion, inputs, labels)

Learning Rate in Reinforcement Learning

In reinforcement learning, the learning rate affects how quickly the agent updates its policy based on new experiences. Here's a simple example using Q-learning.

import numpy as np

def q_learning(env, episodes, alpha, gamma, epsilon):
    Q = np.zeros((env.observation_space.n, env.action_space.n))
    
    for _ in range(episodes):
        state = env.reset()
        done = False
        
        while not done:
            if np.random.random() < epsilon:
                action = env.action_space.sample()
            else:
                action = np.argmax(Q[state])
            
            next_state, reward, done, _ = env.step(action)
            
            # Q-learning update
            Q[state, action] += alpha * (reward + gamma * np.max(Q[next_state]) - Q[state, action])
            
            state = next_state
    
    return Q

# Usage example (not executed)
# env = gym.make('FrozenLake-v1')
# Q = q_learning(env, episodes=10000, alpha=0.1, gamma=0.99, epsilon=0.1)

Additional Resources

For further exploration of learning rates in deep learning, consider these resources:

  1. "Cyclical Learning Rates for Training Neural Networks" by Leslie N. Smith ArXiv: https://arxiv.org/abs/1506.01186
  2. "Bag of Tricks for Image Classification with Convolutional Neural Networks" by Tong He et al. ArXiv: https://arxiv.org/abs/1812.01187
  3. "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour" by Priya Goyal et al. ArXiv: https://arxiv.org/abs/1706.02677
  4. "Fixing Weight Decay Regularization in Adam" by Ilya Loshchilov and Frank Hutter ArXiv: https://arxiv.org/abs/1711.05101

These papers provide in-depth discussions on learning rate techniques and their impact on model performance.

More in Backpropagation and optimization

  1. Forward and Backpropagation in Neural Networks using Python8 min

    Neural Networks are computational models inspired by the human brain, capable of learning from data and making predictions or decisions. They are composed of...

  2. Efficient Optimizers for Neural Network Training11 min

    AdamW and Adam are popular optimization algorithms for training neural networks. These algorithms adapt learning rates for each parameter, which can lead to...

  3. Gradient Clipping and Learning Rate Scheduling in Python.md7 min

    Gradient clipping is a technique used in training deep neural networks to prevent the problem of exploding gradients, where the gradients become too large...

  4. Chain Rule in Neural Network Backpropagation16 min

    The chain rule is a fundamental concept in calculus that plays a crucial role in the backpropagation algorithm used to train neural networks. It allows us to...

  5. Gradient Descent & Optimization Explorer in Python7 min

    Gradient Descent is a fundamental optimization algorithm used in machine learning to minimize the cost function of a model. It iteratively adjusts parameters...

  6. Building Stochastic Gradient Descent from Scratch in Python8 min

    Stochastic Gradient Descent is a fundamental optimization algorithm used in machine learning to minimize the loss function. It's an iterative method that...