Sophia优化器

想将模型训练成本削减50%吗？试试这个全新的即插即用优化器：Sophia

使用方法

使用pip安装 pip install Sophia-Optimizer

import torch 
from torch import nn
from Sophia import SophiaG 

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

#初始化模型、损失函数和输入数据
model = MyModel()
loss_function = nn.CrossEntropy()
input_data = ... #输入数据

#初始化优化器
optimizer = SophiaG(model.parameters(), lr=2e-4, betas=(0.965, 0.99), rho = 0.01, weight_decay=1e-1)

#训练循环
for epoch in range(epochs):
    for batch in data_loader:
        optimizer.zero_grad()
        output = model(batch)
        loss = loss_function(output, target)
        loss.backward()
        optimizer.step()

训练：

使用git克隆方法运行训练

进入experiments文件夹

cd Sophia cd experiments

然后运行文件 python3 training.py

如果不行，请执行以下操作：

from Sophia import DecoupledSophia, trainer

#训练模型
trainer.train()

#评估模型
eval_results = trainer.evaluate()
print(f"困惑度：{torch.exp(torch.tensor(eval_results['eval_loss']))}")

现在experiments文件夹中已准备好训练文件！🔥🔥🔥

Sophia是一种二阶裁剪随机优化算法，它使用Hessian矩阵对角线的廉价随机估计作为预条件，并使用裁剪机制来控制最坏情况下的更新大小。与Adam相比，Sophia在验证预训练损失、总计算量和挂钟时间方面表现更佳。通过将模型训练成本减半，Sophia可以帮助节省数百万甚至数十亿美元的计算资源。

优势

Sophia以50%更少的步数达到与Adam相同的验证预训练损失

总计算量和挂钟时间减少50%

无缝集成到现有训练流程中——即插即用！

对模型架构或计算基础设施没有特殊要求

支持Hutchinson和Gauss-Newton-Bartlett两种Hessian估计器

算法伪代码：

初始化参数：θ1、学习率{ηt}、超参数λ、β1、β2、ϵ和估计器选择Estimator ∈ {Hutchinson, Gauss-Newton-Bartlett}
设置m0 = 0, v0 = 0, h1−k = 0
对于t = 1到T执行：
    计算小批量损失Lt(θt)
    计算gt = ∇Lt(θt)
    mt = β1mt−1 + (1 − β1)gt
    如果t mod k = 1，则：
        计算hˆt = Estimator(θt)
        ht = β2ht−k + (1 − β2)hˆt
    否则：
        ht = ht−1
    θt = θt − ηtλθt（权重衰减）
    θt+1 = θt − ηt · clip(mt/ max{ht, ϵ}, ρ)

PyTorch实现

import torch 

class Sophia(torch.optim.Optimizer):
    def __init__(self, model, input_data, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, k=10, estimator="Hutchinson", rho=1):
        self.model = model
        self.input_data = input_data
        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, k=k, estimator=estimator, rho=rho)
        super(Sophia, self).__init__(params, defaults)

    def step(self, closure=None):
        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None:
                    continue
                grad = p.grad.data
                if grad.is_sparse:
                    raise RuntimeError("Sophia不支持稀疏梯度")
                
                state = self.state[p]

                #状态初始化
                if len(state) == 0:
                    state['step'] = 0
                    state['m'] = torch.zeros_like(p.data)
                    state['h'] = torch.zeros_like(p.data)

                m, h = state['m'], state['h']
                beta1, beta2 = group['betas']
                state['step'] += 1

                if group['weight_decay'] != 0:
                    grad = grad.add(group["weight_decay"], p.data)

                #更新有偏一阶矩估计
                m.mul_(beta1).add_(1 - beta1, grad)

                #更新Hessian估计
                if state['step'] % group['k'] == 1:
                    if group['estimator'] == "Hutchinson":
                        hessian_estimate = self.hutchinson(p, grad)
                    elif group['estimator'] == "Gauss-Newton-Bartlett":
                        hessian_estimate = self.gauss_newton_bartlett(p, grad)
                    else:
                        raise ValueError("无效的估计器选择")
                    h.mul_(beta2).add_(1 - beta2, hessian_estimate)

                #更新参数
                p.data.add_(-group['lr'] * group['weight_decay'], p.data)
                p.data.addcdiv_(-group['lr'], m, h.add(group['eps']).clamp(max=group['rho']))

        return loss
    
    def hutchinson(self, p, grad):
        u = torch.randn_like(grad)
        hessian_vector_product = torch.autograd.grad(grad.dot(u), p, retain_graph=True)[0]
        return u * hessian_vector_product
    
    def gauss_newton_bartlett(self, p, grad):
        B = len(self.input_data)
        logits = [self.model(xb) for xb in self.input_data]
        y_hats = [torch.softmax(logit, dim=0) for logit in logits]
        g_hat = torch.autograd.grad(sum([self.loss_function(logit, y_hat) for logit, y_hat in zip(logits, y_hats)]) / B, p, retain_graph=True)[0]
        return B * g_hat * g_hat

解耦Sophia优化器的超参数调优指南

使用解耦Sophia优化器时，调整超参数对于在特定模型和数据集上获得最佳性能至关重要。以下是解耦Sophia优化器参数调优指南：

学习率（lr）

学习率是一个关键的超参数，它控制优化过程中参数更新的步长。在解耦Sophia中，更新写作p.data.addcdiv_(-group['lr'], m, h.add(group['rho']))，这等同于论文中的更新，只是重新参数化了。

调整学习率的建议：

选择的学习率应约为AdamW使用学习率的一半。一些正在进行的部分结果表明，学习率可能还可以进一步增大，可能导致更快的收敛。

Rho（rho）

rho参数在更新规则中用于控制Hessian对参数更新的影响。选择适当的rho值对于平衡梯度和Hessian信息之间的权衡至关重要。

调整rho的建议：

考虑在0.03到0.04范围内选择rho。rho值似乎可以在不同规模的模型间迁移。例如，rho = 0.03可用于125M和335M的Sophia-G模型。 335M Sophia-G的(lr, rho)选择为(2e-4, 0.03)。虽然我们怀疑学习率可能更大，但对于您的特定用例，尝试不同的值以找到最佳组合是至关重要的。