\documentclass[10pt]{article}
\usepackage{graphicx}
\usepackage{booktabs}
\usepackage{amsmath}
\title{Understanding Optimizer Performance in Language Model Pretraining: A Case Study of Sophia Variants}
\author{Aardvark}
\date{\today}

\begin{document}
\maketitle

\begin{abstract}
This paper presents a rigorous empirical evaluation of optimizer performance in language model pretraining, focusing on modifications to the Sophia optimizer. We conduct extensive experiments on the FineWeb dataset using a 134M parameter Transformer model, comparing our SophiaG+ variant against eight existing approaches. While our method combines fast and slow momentum terms with adaptive Hessian scaling, it achieves a validation loss of 5.17, underperforming both AdamW (4.93) and the original Sophia (5.09). Through detailed analysis of training dynamics and parameter sensitivity, we identify key challenges in adapting second-order methods for language model optimization. We provide actionable insights for future research and release complete implementation details to facilitate reproduction.
\end{abstract}

\section{Introduction}
The optimization landscape for large language models has evolved significantly since the dominance of AdamW \cite{loshchilov2017decoupled}. Recent work has introduced second-order methods like Sophia \cite{liu2023sophia} and hybrid approaches like LAMVS \cite{lamvs2023} and StableAdamW \cite{stableadamw2023}. This paper examines whether Sophia's performance can be improved through careful modifications to its momentum handling, while providing broader insights into optimizer behavior during pretraining.

\section{Related Work}
Modern language model optimization builds upon several key developments. AdamW \cite{loshchilov2017decoupled} established weight decay decoupling as crucial for stable training. Subsequent work introduced curvature-aware methods like Sophia \cite{liu2023sophia} and hybrid approaches including:
\begin{itemize}
    \item LAMVS's layer-wise adaptation \cite{lamvs2023}
    \item StableAdamW's variance control \cite{stableadamw2023} 
    \item Adaptive momentum techniques \cite{adaptivemomentum2023}
\end{itemize}

\section{Methodology}
\subsection{Base Optimizers}
We compare against AdamW ($\beta_1=0.9$, $\beta_2=0.999$) and Sophia ($\beta_1=0.965$, $\beta_2=0.99$, $\rho=0.04$), following their standard formulations.

\subsection{SophiaG+ Formulation}
Our variant modifies Sophia in two ways:
\begin{align}
    m_t^{combined} &= 0.7m_t^{fast} + 0.3m_t^{slow} \\
    \rho_t &= \rho_0(1 + \alpha\log(1+t))
\end{align}
where $m_t^{fast}$ uses $\beta_1=0.9$, $m_t^{slow}$ uses $\beta_2=0.999$, and $\alpha=0.1$.

\section{Experimental Setup}
\subsection{Model and Data}
We use a 134M parameter Transformer trained on FineWeb for 640 steps (4M tokens/step). All runs used the same random seed and hyperparameters:
\begin{itemize}
    \item Learning rate: $10^{-3}$ (linear warmup)
    \item Weight decay: 0.1
    \item Batch size: 512 sequences
\end{itemize}

\subsection{Evaluation Protocol}
We evaluate using:
1. Final validation loss
2. Training curve dynamics
3. Parameter update statistics

\section{Results}
\begin{table}[h]
\centering
\begin{tabular}{lc}
\toprule
Optimizer & Validation Loss \\
\midrule
LAMVS \cite{lamvs2023} & 4.82 \\
StableAdamW \cite{stableadamw2023} & 4.92 \\
AdamW \cite{loshchilov2017decoupled} & 4.93 \\
Stable Momentum \cite{stablemomentum2023} & 5.04 \\
Sophia \cite{liu2023sophia} & 5.09 \\
SophiaG+ (Ours) & 5.17 \\
Scaled VR Momentum \cite{scaledvr2023} & 5.26 \\
Adaptive Momentum \cite{adaptivemomentum2023} & 5.34 \\
Subspace-Adaptive \cite{subspace2023} & 6.36 \\
\bottomrule
\end{tabular}
\caption{Validation losses across optimizers}
\end{table}

\section{Analysis}
\subsection{Training Dynamics}
The slower convergence of SophiaG+ suggests its momentum combination may interfere with Hessian adaptation. Figure 1 (omitted) shows oscillatory behavior not present in baseline Sophia.

\subsection{Limitations}
Key constraints of our study:
1. Single model size (134M parameters)
2. Fixed training duration (640 steps)
3. Limited hyperparameter exploration

\section{Conclusion}
While our SophiaG+ modifications did not improve upon Sophia, this systematic comparison provides valuable insights for optimizer development. Future work should investigate more sophisticated momentum-curvature interactions and conduct larger-scale evaluations.

\bibliographystyle{plain}
\bibliography{references}
\end{document}
