Merge pull request #998 from ChrisRackauckas-Claude/add-sophia-docstring

Add comprehensive docstring for Sophia optimizer

Author: ChrisRackauckas

docs/src/optimization_packages/optimization.md

@@ -8,24 +8,9 @@ There are some solvers that are available in the Optimization.jl package directl
 
     This can also handle arbitrary non-linear constraints through a Augmented Lagrangian method with bounds constraints described in 17.4 of Numerical Optimization by Nocedal and Wright. Thus serving as a general-purpose nonlinear optimization solver available directly in Optimization.jl.
 
-  - `Sophia`: Based on the recent paper https://arxiv.org/abs/2305.14342. It incorporates second order information in the form of the diagonal of the Hessian matrix hence avoiding the need to compute the complete hessian. It has been shown to converge faster than other first order methods such as Adam and SGD.
-    
-      + `solve(problem, Sophia(; η, βs, ϵ, λ, k, ρ))`
-    
-      + `η` is the learning rate
-      + `βs` are the decay of momentums
-      + `ϵ` is the epsilon value
-      + `λ` is the weight decay parameter
-      + `k` is the number of iterations to re-compute the diagonal of the Hessian matrix
-      + `ρ` is the momentum
-      + Defaults:
-        
-          * `η = 0.001`
-          * `βs = (0.9, 0.999)`
-          * `ϵ = 1e-8`
-          * `λ = 0.1`
-          * `k = 10`
-          * `ρ = 0.04`
+```@docs
+Sophia
+```
 
 ## Examples
 

src/sophia.jl

@@ -1,3 +1,49 @@
+"""
+    Sophia(; η = 1e-3, βs = (0.9, 0.999), ϵ = 1e-8, λ = 1e-1, k = 10, ρ = 0.04)
+
+A second-order optimizer that incorporates diagonal Hessian information for faster convergence.
+
+Based on the paper "Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training" 
+(https://arxiv.org/abs/2305.14342). Sophia uses an efficient estimate of the diagonal of the Hessian
+matrix to adaptively adjust the learning rate for each parameter, achieving faster convergence than 
+first-order methods like Adam and SGD while avoiding the computational cost of full second-order methods.
+
+## Arguments
+
+  - `η::Float64 = 1e-3`: Learning rate (step size)
+  - `βs::Tuple{Float64, Float64} = (0.9, 0.999)`: Exponential decay rates for the first moment (β₁) 
+    and diagonal Hessian (β₂) estimates
+  - `ϵ::Float64 = 1e-8`: Small constant for numerical stability
+  - `λ::Float64 = 1e-1`: Weight decay coefficient for L2 regularization
+  - `k::Integer = 10`: Frequency of Hessian diagonal estimation (every k iterations)
+  - `ρ::Float64 = 0.04`: Clipping threshold for the update to maintain stability
+
+## Example
+
+```julia
+using Optimization, OptimizationOptimisers
+
+# Define optimization problem
+rosenbrock(x, p) = (1 - x[1])^2 + 100 * (x[2] - x[1]^2)^2
+x0 = zeros(2)
+optf = OptimizationFunction(rosenbrock, Optimization.AutoZygote())
+prob = OptimizationProblem(optf, x0)
+
+# Solve with Sophia
+sol = solve(prob, Sophia(η=0.01, k=5))
+```
+
+## Notes
+
+Sophia is particularly effective for:
+  - Large-scale optimization problems
+  - Neural network training
+  - Problems where second-order information can significantly improve convergence
+  
+The algorithm maintains computational efficiency by only estimating the diagonal of the Hessian
+matrix using a Hutchinson trace estimator with random vectors, making it more scalable than 
+full second-order methods while still leveraging curvature information.
+"""
 struct Sophia
     η::Float64
     βs::Tuple{Float64, Float64}