class_docstr: ## Sophia-G Optimizer
  
  We extend the class `GenericAdaptiveOptimizer` defined in [`__init__.py`](index.html)
  to implement the Sophia optimizer.
docstr:         ### Update the EMA of Hessian diagonal $h_t$
  
          * `n_tokens_training_batch` is the number of tokens/inputs in the batch $B$
  
          egin{align}
          \hat{h}_t &= B \cdot 
  abla_   heta \hat{L} (  heta) \odot 
  abla_   heta \hat{L} (  heta) \
          h_t &= eta_2 h_{t-k} + (1 - eta_2) \hat{h}_t
          \end{align}
          
code: def update_hessian(self, n_tokens_training_batch):
      """
          ### Update the EMA of Hessian diagonal $h_t$
  
          * `n_tokens_training_batch` is the number of tokens/inputs in the batch $B$
  
          \begin{align}
          \hat{h}_t &= B \cdot \nabla_\theta \hat{L} (\theta) \odot \nabla_\theta \hat{L} (\theta) \\
          h_t &= \beta_2 h_{t-k} + (1 - \beta_2) \hat{h}_t
          \end{align}
          """
  
      # Iterate through parameter groups
      for group in self.param_groups:
          # $\beta_2$
          _, beta2 = group["betas"]
          # Iterate through parameters
          for p in group["params"]:
              # Skip parameters without gradients
              if p.grad is None:
                  continue
  
              # Get optimizer state
              state = self.state[p]
  
              # Initialize state if empty
              if len(state) == 0:
                  self.init_state(state, group, p)
  
              # Update EMA Hessian diagonal
              #
              # \begin{align}
              # \hat{h}_t &= B \cdot \nabla_\theta \hat{L} (\theta) \odot \nabla_\theta \hat{L} (\theta) \\
              # h_t &= \beta_2 h_{t-k} + (1 - \beta_2) \hat{h}_t
              # \end{align}
              state["hessian"].mul_(beta2).addcmul_(
                  p.grad,
                  p.grad,
                  value=(1 - beta2) * n_tokens_training_batch,
              )
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/optimizers/sophia.py
imports: from typing import Dict, Any, Tuple, Optional
  import torch
  from torch import nn
  from labml_nn.optimizers import GenericAdaptiveOptimizer, WeightDecay
label_desc: Use image node_id 3 for calling the Sophia obj's underlying member method code with example usage: obj.update_hessian(n_tokens_training_batch) without return types