This is a PyTorch implementation of Sophia-G from paper Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training. Official implementation is available at Liuhong99/Sophia.
\nSophia is more adaptive to heterogeneous curvatures than Adam, more resistant to non-convexity and rapid change of Hessian than Newton\u2019s method, and also uses a low-cost pre-conditioner.
\nSophia keeps diagonal Hessian estimates with EMA across iterations. The diagonal Hessian _^_0_^_ is calculated every _^_1_^_ steps.
\n_^_2_^_Sophia uses EMA of gradients _^_3_^_, only considers positive entries of the diagonal Hessian and does per-coordinate clipping to the update.
\n_^_4_^_where _^_5_^_ is a very small value to prevent division by _^_6_^_.
\nwhere _^_8_^_ are the inputs, _^_9_^_ is the batch size (number of inputs/tokens), _^_10_^_ is cross entropy loss, and _^_11_^_ are sampled from the logits _^_12_^_.
\nNote that this hessian estimate is always positive and therefore we can replace _^_13_^_ with _^_14_^_.
\nSophia with Gauss-Newton-Bartlett (GNB) estimator is Sophia-G
\nHere is an experiment that uses Sophia-G to train a transformer.
\n": "This is a PyTorch implementation of Sophia-G from paper Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training. Official implementation is available at Liuhong99/Sophia.
\nSophia is more adaptive to heterogeneous curvatures than Adam, more resistant to non-convexity and rapid change of Hessian than Newton\u2019s method, and also uses a low-cost pre-conditioner.
\nSophia keeps diagonal Hessian estimates with EMA across iterations. The diagonal Hessian _^_0_^_ is calculated every _^_1_^_ steps.
\n_^_2_^_Sophia uses EMA of gradients _^_3_^_, only considers positive entries of the diagonal Hessian and does per-coordinate clipping to the update.
\n_^_4_^_where _^_5_^_ is a very small value to prevent division by _^_6_^_.
\nwhere _^_8_^_ are the inputs, _^_9_^_ is the batch size (number of inputs/tokens), _^_10_^_ is cross entropy loss, and _^_11_^_ are sampled from the logits _^_12_^_.
\nNote that this hessian estimate is always positive and therefore we can replace _^_13_^_ with _^_14_^_.
\nSophia with Gauss-Newton-Bartlett (GNB) estimator is Sophia-G
\nHere is an experiment that uses Sophia-G to train a transformer.
\n", "We extend the class _^_0_^_ defined in _^_1_^_ to implement the Sophia optimizer.
\n": "We extend the class _^_0_^_ defined in _^_1_^_ to implement the Sophia optimizer.
\n", "We do the following parameter update,
\n_^_7_^_": "We do the following parameter update,
\n_^_7_^_", "_^_0_^_
\n": "_^_0_^_
\n", "Calculate weight decay
\n": "Calculate weight decay
\n", "Exponential moving average of Hessian diagonal, _^_0_^_
\n": "Exponential moving average of Hessian diagonal, _^_0_^_
\n", "Exponential moving average of gradients, _^_0_^_
\n": "Exponential moving average of gradients, _^_0_^_
\n", "Get _^_0_^_
\n": "Get _^_0_^_
\n", "Get _^_0_^_ and _^_1_^_
\n": "Get _^_0_^_ and _^_1_^_
\n", "Get maximum learning rate _^_0_^_
\n": "Get maximum learning rate _^_0_^_
\n", "Get optimizer state
\n": "Get optimizer state
\n", "In-place calculation of _^_0_^_ _^_1_^_
\n": "In-place calculation of _^_0_^_ _^_1_^_
\n", "Increment _^_0_^_ the number of optimizer steps
\n": "Increment _^_0_^_ the number of optimizer steps
\n", "Initialize state if empty
\n": "Initialize state if empty
\n", "Iterate through parameter groups
\n": "Iterate through parameter groups
\n", "Iterate through parameters
\n": "Iterate through parameters
\n", "Skip parameters without gradients
\n": "Skip parameters without gradients
\n", "This is the number of optimizer steps taken on the parameter, _^_0_^_
\n": "This is the number of optimizer steps taken on the parameter, _^_0_^_
\n", "Update EMA Hessian diagonal
\n_^_0_^_\n": "
Update EMA Hessian diagonal
\n_^_0_^_\n", "A simple PyTorch implementation/tutorial of Sophia optimizer": "A simple PyTorch implementation/tutorial of Sophia optimizer", "Sophia Optimizer": "Sophia Optimizer" }