\documentclass{article}
\usepackage[utf8]{inputenc}
\usepackage{amsmath}
\usepackage{graphicx}
\usepackage{booktabs}

\title{StableAdamW: Variance-Stabilized Optimization for Language Models}

\author{Aardvark}

\date{\today}

\begin{document}

\maketitle

\begin{abstract}
We present StableAdamW, a modified version of AdamW that addresses training instability through controlled variance clipping of second moment estimates. While the performance improvement over AdamW is modest (4.919 vs 4.927 validation loss), our analysis reveals more consistent training dynamics. The method requires no additional memory overhead and maintains the computational efficiency of AdamW.
\end{abstract}

\section{Introduction}
Recent work has shown that adaptive optimization methods like AdamW can exhibit instability during language model training. We hypothesize that much of this instability stems from uncontrolled variance in the second moment estimates.

Our contribution is StableAdamW, which introduces conservative variance clipping of second moment estimates. While our final performance improvement over AdamW is small, the method demonstrates more consistent convergence.

\section{Related Work}
Our work builds upon several research directions in optimization:

\textbf{Adaptive Methods}: Adam and AdamW introduced parameter-specific learning rate adaptation.

\textbf{Variance Reduction}: Various techniques have been proposed to stabilize gradient estimates.

\textbf{Optimizer Stability}: Recent work has identified instability issues in adaptive methods.

\section{Method}
StableAdamW modifies the AdamW update as follows:

\begin{equation}
    \hat{v}_t = \min\left(\sqrt{v_t} + \epsilon, \tau\right)
\end{equation}

where $\tau=2.0$ is the clipping threshold. The update steps are:

1. Compute gradient $g_t$
2. Update first moment $m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t$
3. Update second moment $v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2$
4. Apply variance clipping $\hat{v}_t = \min(\sqrt{v_t} + \epsilon, \tau)$
5. Update parameters $\theta_t = \theta_{t-1} - \eta m_t / \hat{v}_t$

\section{Experimental Setup}
We evaluate on a 134M parameter transformer trained on FineWeb with:
\begin{itemize}
\item Batch size 1024, context length 4096
\item Cosine learning rate schedule
\item Compared against AdamW baseline
\end{itemize}

\section{Results}
\begin{table}[h]
\centering
\begin{tabular}{lc}
\toprule
Method & Validation Loss \\
\midrule
StableAdamW (ours) & 4.919 \\
AdamW & 4.927 \\
\bottomrule
\end{tabular}
\caption{Validation results comparing optimizers.}
\end{table}

\section{Conclusions}
We presented StableAdamW, demonstrating that careful second moment control can improve training stability. The method provides a simple modification to AdamW that may benefit practitioners.

\begin{thebibliography}{10}
\bibitem{adam}
Kingma, Ba. "Adam: A method for stochastic optimization." arXiv 2014.

\bibitem{adamw}
Loshchilov, Hutter. "Decoupled weight decay regularization." arXiv 2017.
\end{thebibliography}

\end{document}
