✅作者简介热爱科研的Matlab仿真开发者擅长数学建模、数据处理、建模仿真、程序设计、完整代码获取、论文复现及科研仿真。 往期回顾关注个人主页Matlab科研工作室 关注我领取海量matlab电子书和数学建模资料个人信条格物致知,完整Matlab代码获取及仿真咨询内容私信。 内容介绍当前深度学习大模型训练场景中业界主流优化器依然以Adam、SGD这类一阶优化器为主流这类方法仅依赖损失函数的一阶梯度信息完成参数更新在训练超大规模大语言模型、计算机视觉大模型时普遍存在收敛速度慢、训练步数多、损失函数在鞍点区域震荡难以跳出的痛点动辄需要数周的训练时间算力成本居高不下。传统基于牛顿法的二阶优化器理论上可以利用损失函数的二阶曲率信息实现更精准的参数更新收敛速度远快于一阶优化器但这类方法需要显式计算和存储海森矩阵单步计算复杂度随模型参数量呈平方级增长完全无法适配参数量达到十亿、千亿级别的现代深度学习大模型长期以来难以实现规模化工业落地。针对上述行业痛点本研究提出二阶剪切随机优化框架设计一款面向深度学习训练的可扩展二阶优化器在几乎不增加单步计算开销的前提下利用随机估计得到的海森矩阵对角线信息构建预条件子通过二阶剪切机制严格控制极端大梯度更新步长的数值范围既保留二阶优化器收敛速度快的核心优势又完全规避传统二阶方法计算复杂度爆炸的问题实现从千万级小模型到千亿级大模型的全场景高效训练大幅降低深度学习模型的整体训练时间与算力成本。相关研究基础与技术缺口近年来深度学习二阶优化方向的相关研究已经形成多个明确的技术分支以L-BFGS为代表的拟牛顿类方法通过历史梯度信息近似拟合海森矩阵逆避免了显式计算海森矩阵但这类方法在深度非凸损失景观中容易出现近似精度快速衰减的问题训练稳定性难以保证以Shampoo为代表的块级二阶优化方法将大模型的参数按张量维度分块在每个子块内执行二阶预条件更新大幅降低了计算复杂度但分块策略需要针对不同模型架构人工定制通用性不足对全连接层、注意力层这类不同结构的层适配性差最新提出的Sophia类二阶优化器通过海森向量积的方式随机估计海森矩阵对角线实现了极低复杂度的二阶信息获取但这类方法缺乏对极端大更新步长的显式约束在大模型训练的早期阶段容易出现参数数值震荡甚至训练发散的问题。当前现有二阶优化方案的共性技术缺口在于没有在统一框架内同时实现通用架构适配的低复杂度二阶信息估计、严格的极端步长剪切机制、可分布式扩展的并行训练能力难以在工业界主流的大规模分布式训练集群中直接落地部署。本研究提出的二阶剪切随机优化可扩展二阶优化器针对性填补上述技术缺口实现全场景通用的高效二阶深度学习训练。二阶剪切随机优化核心理论框架本研究提出的二阶剪切随机优化器Second-order Clipping Stochastic Optimizer, SCSO整体由四大核心模块串联组成所有模块的计算复杂度均与模型参数量呈线性关系完全适配任意参数量级的深度学习模型训练⛳️ 运行结果 部分代码function [p, avg_g, avg_hess] sophiaupdate(p, g, avg_g, avg_hess, hess_est,...update_hess, t, k, lr, beta1, beta2, rho, bs, weight_decay, epsilon)%SOPHIAUPDATE Update parameters via Sophia second-order optimization%% [NET,AVG_G,AVG_H] SOPHIAUPDATE(NET,GRAD,AVG_G,AVG_H,HESS_EST,% HESS_UPDATE,ITER,HESS_ITER) updates the learnable parameters of the% dlnetwork NET using the Sophia (Second-order Clipped Stochastic% Optimization with diagonal Hessian pre-conditioning) algorithm. Sophia% improves upon Adam by adapting step sizes to heterogeneous curvatures% using a diagonal Hessian estimate and per-coordinate clipping.%% Input GRAD contains the gradients of the loss with respect to each of% the network parameters. Inputs AVG_G and AVG_H contain the moving% average of the parameter gradients and the moving average of the% diagonal Hessian estimates, respectively. AVG_G and AVG_H are obtained% from the previous call to SOPHIAUPDATE. GRAD, AVG_G, and AVG_H must be% tables with the same structure as NET.Learnables, with a Value variable% containing a cell array of parameter gradients, average gradients, or% average Hessian estimates. Input GRAD can be obtained using the% dlgradient and dlfeval functions. The global learning rate is% multiplied by the corresponding learning rate factor for each parameter% in each layer of the dlnetwork.%% If inputs AVG_G and AVG_H are empty, the function assumes no previous% gradients and executes like the first update in a series of iterations.%% Input HESS contains the diagonal Hessian estimate for the current step.% Two estimation methods are supported:%% Option 1 - Gauss-Newton-Bartlett (GNB, recommended for LLMs):% Sample labels y_hat ~ Categorical(softmax(logits)), compute the loss% on sampled labels, backpropagate, then pass (grad_sampled .*% grad_sampled) as HESS.%% Option 2 - Hutchinson unbiased estimator:% Draw u ~ N(0,I), compute the Hessian-vector product H*u via a second% backward pass, then pass (u .* Hu) as HESS.%% Pass [] (or zeros) when HESS_UPDATE is false; the stored AVG_H from the% previous estimation step will be reused.%% Input HESS_UPDATE is a logical flag. Set HESS_UPDATE to TRUE on Hessian% estimation steps (every k iterations, e.g. k10), FALSE otherwise.%% Input ITER contains the update iteration number. ITER must be a% positive integer. Use a value of 1 for the first call to SOPHIAUPDATE% and increment by 1 for each successive call in a series of iterations.% The Sophia algorithm uses this value to correct for bias in the moving% averages at the beginning of a set of iterations.%% Input HESS_ITER contains the update iteration number of the Hessian.%% Outputs NET, AVG_G, and AVG_H are the updated dlnetwork, average% gradients, and average Hessian estimates, respectively.%% [PARAMS,AVG_G,AVG_H] SOPHIAUPDATE(PARAMS,GRAD,AVG_G,AVG_H,HESS,% HESS_UPDATE,ITER,HESS_ITER) updates the deep learning parameters in% PARAMS using the Sophia algorithm. Input PARAMS can be a dlarray, a% numeric array, a cell array, a structure, or a table with a Value% variable containing the learnable parameters of the network. GRAD,% AVG_G, and AVG_H must have the same datatype and ordering as PARAMS.% Input GRAD can be obtained using the dlgradient and dlfeval functions.% All parameter values are updated using the global learning rate.%% Outputs PARAMS, AVG_G, and AVG_H are the updated parameters, average% gradients, and average Hessian estimates, respectively.%% [___] SOPHIAUPDATE(___,LEARNRATE,BETA1,BETA2,RHO,BATCHSIZE,% WEIGHT_DECAY) also specifies values to use for the global learning% rate, gradient EMA decay factor, Hessian EMA decay factor, effective% batch size in tokens, per-coordinate clipping threshold, and decoupled% weight decay factor. LEARNRATE must be a positive scalar. BETA1, BETA2,% and RHO must be scalars between 0 and 1. BATCHSIZE and WEIGHT_DECAY% must be positive scalars.%% [___] SOPHIAUPDATE(___,LEARNRATE,BETA1,BETA2,RHO,BATCHSIZE,% WEIGHT_DECAY,EPSILON) specifies a small constant used to prevent% division by zero in the update equation. The default value of EPSILON% is 1e-12.%% Default values:% LEARNRATE 1e-4% BETA1 0.9 (gradient EMA decay)% BETA2 0.99 (Hessian EMA decay)% RHO 0.01 (clipping threshold)% BATCHSIZE 1 (tokens per update)% WEIGHT_DECAY 0.0 (decoupled weight decay)% EPSILON 1e-12 (numerical stability)%% Example:% % Perform Sophia updates with Hessian estimation every 10 steps.% p rand(100,100);% avg_g [];% avg_h [];% hess_interval 10;% lr 3e-4;%% iter_hess 0;% for iter 1:50% do_hess (iter 1) || (mod(iter, hess_interval) 0);%% % Compute gradients and Hessian estimate (if needed)% [loss, grad] dlfeval(modelGradients, model, X, Y);%% hess [];% if do_hess% iter_hess iter_hess 1;%% [~, grad_sampled] dlfeval(sampledGradients, model, X);% hess 480 * 1024 .* grad_sampled .* grad_sampled;% end%% % Update parameters% [p, avg_g, avg_h] sophiaupdate(p, grad, avg_g, avg_h, ...% hess, do_hess, iter, iter_hess, lr, 0.965, 0.99, 0.04, ...% 491520, 0.1, 1e-15);% end%% Reference:% Liu et al. (2023). Sophia: A Scalable Stochastic Second-order% Optimizer via Non-Convex Cubic Regularization. arXiv:2305.14342.%% See also ADAMUPDATE, DLGRADIENT, DLFEVAL, DLNETWORK, DLUPDATEargumentspgavg_gavg_hesshess_estupdate_hess (1,1) logicalt (1,1) {mustBeNumeric, mustBePositive, mustBeInteger}k (1,1) {mustBeNumeric, mustBePositive, mustBeInteger} tlr (1,1) {mustBeNumeric, mustBeFinite, mustBeNonnegative} 1e-4;beta1 (1,1) {mustBeNumeric, mustBeGreaterThanOrEqual(beta1,0), mustBeLessThan(beta1,1)} 0.9;beta2 (1,1) {mustBeNumeric, mustBeGreaterThanOrEqual(beta2,0), mustBeLessThan(beta2,1)} 0.99;rho (1,1) {mustBeNumeric, mustBePositive} 0.01;bs (1,1) {mustBeNumeric, mustBePositive} 1;weight_decay (1,1) {mustBeNumeric, mustBeFinite, mustBeNonnegative} 0.0;epsilon (1,1) {mustBeNumeric, mustBeFinite, mustBePositive} 1e-12;end% Ensure inputs are dlarray where necessaryif isnumeric(g) isdlarray(p)g dlarray(g);end% 1. Buffer Initialization (Stateless container expansion)% If buffers are empty, we initialize them to the same structure as gif isempty(avg_g)avg_g dlupdate((x) zeros(size(x), like, x), g);endif isempty(avg_hess)avg_hess dlupdate((x) -ones(size(x), like, x), g);endif isempty(hess_est)if update_hesserror(sophiaupdate:HESS_EST cannot be empty when HESS_UPDATE is true.);end% Create a dummy container of zeros so paramArgs structure remains constanthess_est dlupdate((x) zeros(size(x), like, x), g);end% 2. Setup Container Mappingpersistent funcif isempty(func)func deep.internal.LearnableUpdateFunction( ...iSingleStepValue, ...iSingleStepParameter);end% All four inputs are now guaranteed to be containers of the same structureparamArgs {g, matlab.lang.internal.move(avg_g), matlab.lang.internal.move(avg_hess), hess_est};fixedArgs {update_hess, t, k, lr, beta1, beta2, rho, bs, weight_decay, epsilon};% Apply update over network/structure[p, ~, avg_g, avg_hess, ~] deep.internal.networkContainerFixedArgsFun( ...func, p, matlab.lang.internal.move(paramArgs), fixedArgs);endfunction [p, g, avg_g, avg_hess, hess_est] iSingleStepParameter(p, g, avg_g, avg_hess, ...hess_est, update_hess, t, k, lr, beta1, beta2, rho, bs, weight_decay, epsilon)% Update logic for dlnetwork parameters (Learnables table)% Apply per-parameter learn-rate factorlr lr * p.LearnRateFactor;% Apply Weight Decayp.Value p.Value .* (1 - lr * weight_decay);% Apply a correction factor due to the trailing averages being biased% towards zero at the beginning. This is fed into the clipping threshold.biasCorrection (1-beta1.^t)./(1-beta2.^k);effectiveRho biasCorrection.*rho;v p.Value;p.Value []; % Clear to save memory during internal update[v, avg_g, avg_hess] internal_sophia(v, g, avg_g, avg_hess, hess_est, ...lr, beta1, beta2, effectiveRho, bs, epsilon, update_hess);p.Value v;endfunction [p, g, avg_g, avg_hess, hess_est] iSingleStepValue(p, g, avg_g, avg_hess, ...hess_est, update_hess, t, k, lr, beta1, beta2, rho, bs, weight_decay, epsilon)% Update logic for raw dlarrays or numeric arrays% Apply Weight Decay using base LRp p .* (1 - lr * weight_decay);% Apply a correction factor due to the trailing averages being biased% towards zero at the beginning. This is fed into the clipping threshold.biasCorrection (1-beta1.^t)./(1-beta2.^k);effectiveRho biasCorrection.*rho;[p, avg_g, avg_hess] internal_sophia(p, g, avg_g, avg_hess, hess_est, ...lr, beta1, beta2, effectiveRho, bs, epsilon, update_hess);endfunction [l, avg_g, avg_hess] internal_sophia(l, g, avg_g, avg_hess, hess_est, ...lr, beta1, beta2, rho, bs, epsilon, update_hess)%INTERNAL_SOPHIA Internal helper that applies one Sophia optimizer step.%% Applies one complete Sophia update step, delegating to sophiastep for% the gradient-EMA / Hessian-EMA / clipping logic.% Matches the internal_adam.m pattern: skips update when LEARNRATE 0.if lr ~ 0[step, avg_g, avg_hess] sophiastep(g, avg_g, avg_hess, hess_est, ...lr, beta1, beta2, rho, bs, epsilon, update_hess);l l step; % Add the step for Gradient Descentendendfunction [step, avg_g, avg_hess] sophiastep(g, avg_g, avg_hess, hess_est, ...learnrate, beta1, beta2, rho, bs, epsilon, update_hess)%SOPHIASTEP Calculate Sophia update step for a single parameter tensor.%% [STEP,AVG_G,AVG_H] SOPHIASTEP(GRAD,AVG_G,AVG_H,HESS,LEARNRATE,% BETA1,BETA2,RHO,BATCHSIZE,EPSILON,UPDATE_H) computes the parameter% update step using the Sophia algorithm. Sophia uses a diagonal% Hessian-based pre-conditioner with per-coordinate clipping to adapt% step sizes to heterogeneous curvatures across parameter dimensions.%% Input GRAD is the current gradient. Inputs AVG_G and AVG_H contain the% running exponential moving averages of gradients and diagonal Hessian% estimates, respectively. These are updated and returned for use in the% next call.%% Input HESS is the new diagonal Hessian estimate for the current step.% Pass [] or zeros when UPDATE_H is false; the stored AVG_H from the% previous estimation step will be reused.%% Input LEARNRATE is the effective learning rate after bias correction.% Input BETA1 is the gradient EMA decay factor. Input BETA2 is the Hessian% EMA decay factor. Input RHO is the per-coordinate clipping threshold.% Input BATCHSIZE is the effective batch size (batch_size *% sequence_length). Input EPSILON is a small constant for numerical% stability. Input UPDATE_H is a logical flag; set to TRUE to update the% Hessian EMA with HESS, FALSE otherwise.%% Outputs STEP is the computed parameter update. Outputs AVG_G and AVG_H% are the updated gradient and Hessian EMAs, respectively.%% See also SOPHIAUPDATE, ADAMSTEP% 1. Gradient EMAavg_g beta1.*avg_g (1 - beta1).*g;% 2. Hessian EMAif update_hess% Use unbiased start trick: if 0, start at current estimateisFirst all(avg_hess 0);avg_hess isFirst.*hess_est (1 - isFirst).*(beta2.*avg_hess (1 - beta2).*hess_est);end% 3. Sophia Stepdenom rho.*bs.*max(avg_hess, 0) epsilon;ratio min(abs(avg_g)./denom, 1);step -learnrate.*sign(avg_g).*ratio;end 参考文献Liu et al. (2023). *Sophia: A Scalable Stochastic Second-order Optimizer for LanguageModel Pre-training.* [arXiv:2305.14342](https://arxiv.org/abs/2305.14342)往期回顾扫扫下方二维码