PIDAO — A new way of Deep Learning Training Optimization
Combining Dynamic Control with Deep Learning
Understanding usage of Proportional Integral and Derviate Controller for Deep Learning Training Optimization
Today, we will try to understand and implement the paper above. The paper can be downloaded here https://www.nature.com/articles/s41467-024-54451-3.
Why? Because learning how to understand paper and implement it will increase my expertise on ML and AI. And PID is something that I learn extensively during my university years as Process Engineer, so it’s good that I can finally combine both of these worlds.
First, let’s extract the abstract of the paper:
High-performance optimization algorithms are essential in deep learning. However, understanding the behavior of optimization (i.e., learning process) remains challenging due to the instability and weak interpretability of the algorithms. Since gradient-based optimizations can be interpreted as continuous-time dynamical systems, applying feedback control to the dynamical systems that model the optimizers may provide another perspective for exploring more robust, accurate, and explainable optimization algorithms. In this study, we present a framework for optimization called the controlled heavy-ball optimizer. By employing the proportional-integral-derivative (PID) controller in the optimizer, we develop a deterministic continuous-time optimizer called Proportional-Integral-Derivative Accelerated Optimizer (PIDAO) and provide theoretical convergence analysis of PIDAO in unconstrained (non-)convex optimizations. As a byproduct, we derive PIDAO-family schemes for training deep neural networks using specific discretization methods. Compared to classical optimizers, PIDAO empirically demonstrates a more aggressive capacity to explore the loss landscape with lower computational costs due to the properties of the PID controller. Experimental evaluations demonstrate that PIDAO can accelerate convergence and enhance the accuracy of deep learning, achieving state-of-the-art performance compared to advanced algorithms.
Ok, that’s a lot of buzz words. It is very hard to understand it. But let’s do this sections by sections
High-performance optimization algorithms are essential in deep learning. However, understanding the behavior of optimization (i.e., learning process) remains challenging due to the instability and weak interpretability of the algorithms
This part here is the main issues we have with current method. For all of you who are not familiar with optimization techniques used in the deep learning, let me introduce you to few of thems. We have the famous one, which is Adam (Adapative Moment Estimation), the first one that I learn, SGD (Stochastic Gradient Descent), Adagrad (Adaptive Gradient Algorithm) and RMSProp (Root mean square propagation). All of these traditional optimizer can be called through pytorch optimizer. The paper states that these techniques have a lot of issues, among them is, instability and weak interpretability.
Since gradient-based optimizations can be interpreted as continuous-time dynamical systems, applying feedback control to the dynamical systems that model the optimizers may provide another perspective for exploring more robust, accurate, and explainable optimization algorithms
So, in this section, they’re proposing us to use a feedback control in the dynamic systems (the optimizer) to allows us to interprete the optimization process better. This is something that I can understasnd. As process engineer, we have learned about dynamic systems (which is our refinery/process plant) and the feedback control needed to sustain it’s operation. And the feedback system that we use in our plant is PID control (Proportional Integral and Derivative Control). That’s why in the next section they propose:
In this study, we present a framework for optimization called the controlled heavy-ball optimizer. By employing the proportional-integral-derivative (PID) controller in the optimizer, we develop a deterministic continuous-time optimizer called Proportional-Integral-Derivative Accelerated Optimizer (PIDAO) and provide theoretical convergence analysis of PIDAO in unconstrained (non-)convex optimizations
The keyword here is the word deterministic and discrete. By using PIDAO, the optimizer becomes highly interpretable. So, it becomes easily explained and easier to understand leading to better perfromance in non-convex optimization. Non-convex means that the loss landscape is not like the slope of a hill. It has sudden peaks and slopes leading to harder optimization by traditional optimizer.
As a byproduct, we derive PIDAO-family schemes for training deep neural networks using specific discretization methods. Compared to classical optimizers, PIDAO empirically demonstrates a more aggressive capacity to explore the loss landscape with lower computational costs due to the properties of the PID controller. Experimental evaluations demonstrate that PIDAO can accelerate convergence and enhance the accuracy of deep learning, achieving state-of-the-art performance compared to advanced algorithms
I just combine the last section because it basically list down the benefit of using PIDAO. It is more explainable, have more aggressive capavity, lower compatutional cost and can accelertate convergence and enhance the accuracy of deep learning. So int theory it has better performance compared to advanced algorithms. That’s a lot of claims. Let’s see whether we can replicate their claims
Understanding PID
But first, let me put to you here what I understand about PID, the algorithm that I have learned extensively during my university years.
So there are three parts of the algorithm, and as you guessed it, it is the proportional, integral and the derivative. The algorithm is written as below altogether, but each with it’s own sections.
Because the algorithm is so simple, that’s why it’s been use extensively as a superior feedback control algoritm in the oil and gas industry and in the manufacturing plants. And now, they’re bringing this simple but powerful algorithm to the deep learning and AI.
What I Understand
Alright, now we go to the part where I understand. I know that PID is a very simple algorithm and use extensively in the plant. And I also know that current deep learning optimization method is not as simple as PID. Why? Let’s take Adam for example. As I mentioned above, Adam stands for adaptive moment estimation. What it did was, it adapt to the momentum or the gradient by keeping the moving average of it. It did this through statistics, by keeping tracks of all the moving averages of the gradients. So, that’s why it uses a lot of computaiton powers. It needs to keep track to all of that moving averages.
PID will keep this simpler, by only focusing on the current value of the gradient. Yes,it does takes into account the error and and the past mistakes, but it did that through the current data available, not through statistics. So, it will uses less computational power, and it keep it simple so that people can track what actually happened during training
So, that’s basically my understanding.But, we do have a lot of more parts to understand, let’s continue
Representation of Deep Learning Training Process
The imperfect picture that I generate with AI above is how they represent the training process traditionally. Normally, they represent it as a pyhsical model, in this case, a hill slope. And our training, is like a ball going down the slope. We have to control the movement of the ball to ensure that it reach it’s endpoint, which is the lowest minimum point. But, in reality not all hill slope has a smooth curve downward. Sometimes, there’s another slope in the middle that may cause our ball to stop and does not reach it’s lowest minimum point. The optimizer was supposed to solve this issues. By controlling the momentum of the ball, we will ensure that it will not be stuck in the middle and reach the lowest point. That’s why the current optimizer this pyhsical algorithm (momentum etc) to optimize the training process.
This paper propose a paradigm shift eventhough it’s not the first in doing so. Instead of a pyhsical model, it envisions the training process as a continous time dynamic system. And since it’s time-dynamic we can use feedback control to control the training process. That’s where PID comes.
Let’s implement the PIDAO (Proportional Integral Derivative Accelerated Optimizer)
In the PIDAO paper, they recommend that we can control the gradient of training using the PID algorithm. We can just replace the error in the PID to our gradient. Therefore, the PID formula can be modified as below:
The formula for the control input ( u(t) ) in the PIDAO optimizer is:
This paper doesn’t ask us to completely replace the current methods. It recommends us to combine the traditional physical model optimization (the Adam and classical momentum model) with the PID method. We still assume it as pyhsical, but envisioned it as heavy rolling ball down the slope, and we look it as a dynamic continous movement through the slope. So with that, we need to add our PID to the continous-time heavy-ball dynamics. Therefore the PIDAO framework can be described as:
So, let’s implement this PIDAO, let’s implement the basic version and test it with different case. The code of implementation are as below:
#let's implement this class
import torch
from torch.optim.optimizer import Optimizer
class PIDAO(Optimizer):
def __init__(self,params, lr=1e-3, kp =1, ki=0.1, kd=0.01, a=0.1):
"""
PIDAO (Proportional Integral Derivative Accelerated Optimizer
This optimizer includes PID algorithm on top of classical physical
model optimizer. Based on the paper by several scientists from
China
Args:
params: Iterable of parameters to optimize
lr = learning rate
kp: Proportional gain
ki = integral gain
kd: derivative gain
a: Damping coefficient
"""
defaults = dict(lr=lr, kp=kp, ki=ki, kd=kd, a=a)
super(PIDAO,self).__init__(params,defaults)
@torch.no_grad()
def step(self,closure = None):
"""
Performs a single optimization step using PIDAO optimizer for each time this function was called
Args:
closure: A closure that re-evaluates the model and return the loss
"""
#1. update the loss function
loss = None #initialize the loss
if closure is not None:
with torch.enable_grad():
loss = closure() #update the loss based on the closure
#2. update the parameters
for group in self.param_groups:
for p in group['params']:
if p.grad is None:
continue
grad = p.grad.data
state = self.state[p]
#3. state initialization
if len(state) == 0:
state['velocity'] = torch.zeros_like(p.data)
state['integral'] = torch.zeros_like(p.data)
state['prev_grad'] = torch.zeros_like(p.data)
#4. retrieve the state variables
velocity = state['velocity']
integral = state['integral']
prev_grad = state['prev_grad']
#5. get the hyperparameters
lr = group['lr']
kp = group['kp']
ki = group['ki']
kd = group['kd']
a = group['a']
#6. update all the PID parameters
integral += grad
derivative = grad - prev_grad
# PID control update
update = kp * grad + ki * integral + kd * derivative
# Velocity update (momentum + damping)
velocity = a * velocity - lr * update
# Parameter update
p.data += velocity
# Save state
state['velocity'] = velocity
state['integral'] = integral
state['prev_grad'] = grad
return lossAnd if we test it for MNIS handwritten recognition using simple NN, we can get below results.
import torch
from torch.optim.optimizer import Optimizer
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
# Simple Neural Network for MNIST
class SimpleNN(nn.Module):
def __init__(self):
super(SimpleNN, self).__init__()
self.fc1 = nn.Linear(28 * 28, 128)
self.fc2 = nn.Linear(128, 64)
self.fc3 = nn.Linear(64, 10)
def forward(self, x):
x = x.view(-1, 28 * 28)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
# MNIST Test Case
if __name__ == "__main__":
# Transform and DataLoader
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
train_dataset = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
# Model, Loss, Optimizers
model_pidao = SimpleNN()
model_adam = SimpleNN()
criterion = nn.CrossEntropyLoss()
optimizer_pidao = PIDAO(model_pidao.parameters(), lr=0.01, kp=1.0, ki=0.1, kd=0.01, a=0.1)
optimizer_adam = torch.optim.Adam(model_adam.parameters(), lr=0.01)
# Track Losses
pidao_losses = []
adam_losses = []
# Training Loop
for epoch in range(10): # Train for 10 epochs
model_pidao.train()
model_adam.train()
epoch_loss_pidao = 0.0
epoch_loss_adam = 0.0
for batch_idx, (data, target) in enumerate(train_loader):
# PIDAO Optimization
optimizer_pidao.zero_grad()
output_pidao = model_pidao(data)
loss_pidao = criterion(output_pidao, target)
loss_pidao.backward()
optimizer_pidao.step()
epoch_loss_pidao += loss_pidao.item()
# Adam Optimization
optimizer_adam.zero_grad()
output_adam = model_adam(data)
loss_adam = criterion(output_adam, target)
loss_adam.backward()
optimizer_adam.step()
epoch_loss_adam += loss_adam.item()
# Store Average Loss for Epoch
pidao_losses.append(epoch_loss_pidao / len(train_loader))
adam_losses.append(epoch_loss_adam / len(train_loader))
print(f"Epoch {epoch + 1}: PIDAO Loss = {pidao_losses[-1]:.4f}, Adam Loss = {adam_losses[-1]:.4f}")
# Visualization
plt.figure(figsize=(10, 5))
plt.plot(range(1, 11), pidao_losses, label="PIDAO Loss")
plt.plot(range(1, 11), adam_losses, label="Adam Loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.title("Loss Comparison: PIDAO vs Adam")
plt.legend()
plt.show()
Oops, that’s not a very good results. We can see that PIDAO loss increased exponentially. Does this mean their claim is refuted? But based on the paper, PIDAO (the base version) is good for several case. Such as:
1. Quadratic Loss
We can see that PIDAO achieved convergence faster than Adam.
2. RNN with Longer Sequences
And yes, RNN with Longer sequences certainly benefited from using PIDAO vs Adam.
3. Oscillatory Loss
This does seem like not a good comparison, But we can see base version of PIDAO already performs almost as good as Adam.
Recommendations from the Paper
The papers seems to acknowledge that base PIDAO will not generalized well for all optimization tasks, so they recommend several added algorithms. I wouln’t cover them all since it will be too long. But one of the recommendations is by using adaptive version of the PIDAO optimizer with Root Mean Square (RMS) scaling.
Let us implement it by using the code from their GitHub:
https://github.com/NoulliCHEN/PIDAO/tree/main
import torch
from torch.optim.optimizer import Optimizer, required
class PIDAccOptimizer_SI_AAdRMS(Optimizer):
r"""Implements stochastic gradient descent (optionally with momentum).
Nesterov momentum is based on the formula from
`On the importance of initialization and momentum in deep learning`__.
Args:
params (iterable): iterable of parameters to optimize or dicts defining
parameter groups
lr (float): learning rate
momentum (float, optional): momentum factor (default: 0)
weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
dampening (float, optional): dampening for momentum (default: 0)
nesterov (bool, optional): enables Nesterov momentum (default: False)
Example:
>>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
>>> optimizer.zero_grad()
>>> loss_fn(model(input), target).backward()
>>> optimizer.step()
__ http://www.cs.toronto.edu/%7Ehinton/absps/momentum.pdf
.. note::
The implementation of SGD with Momentum/Nesterov subtly differs from
Sutskever et. al. and implementations in some other frameworks.
Considering the specific case of Momentum, the update can be written as
.. math::
v = \rho * v + g \\
p = p - lr * v
where p, g, v and :math:`\rho` denote the parameters, gradient,
velocity, and momentum respectively.
This is in contrast to Sutskever et. al. and
other frameworks which employ an update of the form
.. math::
v = \rho * v + lr * g \\
p = p - v
The Nesterov version is analogously modified.
"""
def __init__(self, params, lr=required, beta1=0.999, beta2=0.9, eps=1e-8, momentum=0.1, dampening=0,
weight_decay=0, nesterov=False, kp=5., ki=0.4, kd=8.):
defaults = dict(lr=lr, momentum=momentum, dampening=dampening, beta1=beta1, beta2=beta2, eps=eps,
weight_decay=weight_decay, nesterov=nesterov, kp=kp, ki=ki, kd=kd)
if nesterov and (momentum <= 0 or dampening != 0):
raise ValueError("Nesterov momentum requires a momentum and zero dampening")
super(PIDAccOptimizer_SI_AAdRMS, self).__init__(params, defaults)
self.k = 1
def __setstate__(self, state):
super(PIDAccOptimizer_SI_AAdRMS, self).__setstate__(state)
for group in self.param_groups:
group.setdefault('nesterov', False)
def step(self, closure=None):
"""Performs a single optimization step.
Arguments:
closure (callable, optional): A closure that reevaluates the model
and returns the loss.
"""
loss = None
if closure is not None:
loss = closure()
for group in self.param_groups:
weight_decay = group['weight_decay']
momentum = group['momentum']
dampening = group['dampening']
nesterov = group['nesterov']
lr = group['lr']
kp = group['kp']
ki = group['ki']
kd = group['kd']
beta1 = group['beta1']
beta2 = group['beta2']
eps = group['eps']
for p in group['params']:
if p.grad is None:
continue
d_p = p.grad.data
if weight_decay != 0:
d_p.add_(p.data, alpha=weight_decay)
if momentum != 0:
param_state = self.state[p]
if 'square_avg' not in param_state:
square_avg = param_state['square_avg'] = torch.zeros_like(p.data)
square_avg.mul_(beta1).addcmul_(d_p, d_p, value=1 - beta1)
else:
square_avg = param_state['square_avg']
square_avg.mul_(beta1).addcmul_(d_p, d_p.conj(), value=1 - beta1)
self.k += 1
avg = square_avg.clone().detach().mul_((1 - beta1 ** 2) ** -1).sqrt().add_(eps)
if 'z_buffer' not in param_state:
z_buf = param_state['z_buffer'] = torch.zeros_like(p.data)
z_buf.add_(d_p, alpha=lr)
else:
z_buf = param_state['z_buffer']
z_buf.add_(d_p, alpha=lr)
correct_z_buf = z_buf.clone().detach().div_(avg)
# correct_z_buf = z_buf.clone().detach()
if 'y_buffer' not in param_state:
param_state['y_buffer'] = torch.zeros_like(p.data)
y_buf = param_state['y_buffer']
# y_buf.addcdiv_(d_p, avg, value=-lr * (kp - momentum * kd)).\
# add_(z_buf, alpha=-ki * lr)
y_buf.addcdiv_(d_p, avg, value=-lr * (kp - momentum * kd)). \
add_(correct_z_buf, alpha=-ki * lr)
y_buf.mul_((1 + momentum * lr) ** -1)
else:
y_buf = param_state['y_buffer']
y_buf.addcdiv_(d_p, avg, value=-lr * (kp - momentum * kd)). \
add_(correct_z_buf, alpha=-ki * lr)
y_buf.mul_((1 + momentum * lr) ** -1)
d_p = torch.zeros_like(p.data).add_(y_buf, alpha=lr).addcdiv_(d_p, avg, value=-kd * lr)
p.data.add_(d_p)
return lossSo, if we test this optimizer for MNIST handwritten recognition tasks as below, we will get this results
import torch
from torch.optim.optimizer import Optimizer
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
# Define the MNIST model
class SimpleNN(nn.Module):
def __init__(self):
super(SimpleNN, self).__init__()
self.fc1 = nn.Linear(28 * 28, 128)
self.fc2 = nn.Linear(128, 64)
self.fc3 = nn.Linear(64, 10)
def forward(self, x):
x = x.view(-1, 28 * 28)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
# Define the PIDAccOptimizer_SI_AAdRMS (provided by user)
# Testing PIDAccOptimizer_SI_AAdRMS vs Adam on MNIST
if __name__ == "__main__":
# MNIST Dataset
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
train_dataset = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
# Define models, criterion, and optimizers
model_pid = SimpleNN()
model_adam = SimpleNN()
criterion = nn.CrossEntropyLoss()
optimizer_pid = PIDAccOptimizer_SI_AAdRMS(
model_pid.parameters(),
lr=0.001, # Lower learning rate for stability
beta1=0.999,
beta2=0.9,
eps=1e-8,
momentum=0.1,
kp=5.0,
ki=0.4,
kd=8.0
)
optimizer_adam = torch.optim.Adam(model_adam.parameters(), lr=0.001)
# Track losses
pid_losses = []
adam_losses = []
# Training loop
for epoch in range(10): # Train for 10 epochs
model_pid.train()
model_adam.train()
epoch_loss_pid = 0.0
epoch_loss_adam = 0.0
for batch_idx, (data, target) in enumerate(train_loader):
# PID Optimization
optimizer_pid.zero_grad()
output_pid = model_pid(data)
loss_pid = criterion(output_pid, target)
loss_pid.backward()
optimizer_pid.step()
epoch_loss_pid += loss_pid.item()
# Adam Optimization
optimizer_adam.zero_grad()
output_adam = model_adam(data)
loss_adam = criterion(output_adam, target)
loss_adam.backward()
optimizer_adam.step()
epoch_loss_adam += loss_adam.item()
# Store average loss for each epoch
pid_losses.append(epoch_loss_pid / len(train_loader))
adam_losses.append(epoch_loss_adam / len(train_loader))
print(f"Epoch {epoch + 1}: PID Loss = {pid_losses[-1]:.4f}, Adam Loss = {adam_losses[-1]:.4f}")
# Visualization
plt.figure(figsize=(10, 5))
plt.plot(range(1, 11), pid_losses, label="PIDAccOptimizer_SI_AAdRMS Loss")
plt.plot(range(1, 11), adam_losses, label="Adam Loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.title("Loss Comparison: PIDAccOptimizer_SI_AAdRMS vs Adam")
plt.legend()
plt.show()
Not bad huh? Seems like PIDAO can be seen to generalize well with multiple task. Now it seems promising as replacement for Adam.
Conclusion:
So what have we learn today? We learn about PID algorithm and how it’s being extensively used in manufacturing plants and in oil and gas industry. We also know what is the current method for optimization in deep learning training.
And we have succesfully implemented PIDAO algorithm and test it for different case. It seems like it generalized well and can be used as alternative/replacement to Adam and other optimization method.
Alright,that’s all for today. Bye