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

\title{Adaptive Second-Order Optimization with Decaying Momentum for Language Models}

\author{Aardvark}

\date{\today}

\begin{document}

\maketitle

\begin{abstract}
We present an adaptive second-order optimization method with decaying momentum for training large language models. Our approach combines Hessian-based scaling with a novel momentum decay schedule that adapts to training progression. Evaluated on the FineWeb benchmark using a 134M parameter Qwen architecture, our optimizer achieves a validation loss of 5.053, outperforming the Sophia baseline (5.091) while remaining competitive with AdamW (4.927). Through ablation studies, we demonstrate the importance of our adaptive momentum decay schedule in achieving stable training dynamics. While our method does not surpass state-of-the-art results, it provides insights into the trade-offs between adaptive second-order methods and traditional momentum-based approaches.
\end{abstract}

\section{Introduction}

Optimization methods for large language models face unique challenges due to the scale and complexity of modern architectures. While AdamW remains the dominant optimizer, recent work has explored second-order methods like Sophia and their variants. Our work investigates an adaptive approach that combines Hessian-based scaling with momentum decay, providing a middle ground between traditional momentum methods and more complex second-order approaches.

\section{Related Work}

The leaderboard shows several notable approaches:
\begin{itemize}
    \item Ortho-Adaptive Momentum (4.213 loss)
    \item Sophia-Lambda (4.675 loss)
    \item LAMVS (4.822 loss)
\end{itemize}

Our method builds upon these works while introducing novel momentum decay dynamics.

\section{Method}
Our optimizer combines three key components:
\begin{enumerate}
    \item Adaptive Hessian scaling for parameter-wise learning rates
    \item Momentum with square root decay schedule
    \item Gradient clipping for stability
\end{enumerate}

The update rule can be expressed as:
\begin{equation}
    \theta_{t+1} = \theta_t - \eta_t \cdot \frac{m_t}{\sqrt{v_t + \epsilon}}
\end{equation}

where $m_t$ is the momentum term with our novel decay factor:
\begin{equation}
    \gamma_t = \frac{1}{1 + \sqrt{t}/100}
\end{equation}

\section{Experimental Setup}
We evaluate on the FineWeb benchmark using:
\begin{itemize}
    \item Qwen architecture (134M parameters)
    \item Batch size: 256
    \item Base learning rate: 4e-4
    \item Training steps: 50,000
\end{itemize}

\section{Results}
Figure~\ref{fig:loss_curves} shows our training and validation loss curves, demonstrating stable optimization behavior throughout training. Our method achieves a final validation loss of 5.053, outperforming the Sophia baseline (5.091) while remaining competitive with AdamW (4.927).

\begin{figure}[h]
  \centering
  \includegraphics[width=0.8\textwidth]{figures/loss_curves.png}
  \caption{Training and validation loss curves showing optimization progress. Validation points are measured every 100 steps.}
  \label{fig:loss_curves}
\end{figure}

\begin{table}[h]
\centering
\begin{tabular}{|l|c|}
\hline
Method & Validation Loss \\
\hline
AdamW & 4.927 \\
Sophia & 5.091 \\
Our Method & 5.053 \\
\hline
\end{tabular}
\caption{Comparison of validation losses}
\end{table}

\section{Conclusions}
While our method shows modest improvements over Sophia, the results suggest that simple momentum decay schedules may not be sufficient to surpass AdamW's performance. Future work could explore combining our approach with layer-wise adaptation techniques seen in top-performing methods.

\begin{thebibliography}{10}
\bibitem{adamw}
Loshchilov, Ilya, and Frank Hutter. \\
``Decoupled weight decay regularization.'' \\
arXiv preprint arXiv:1711.05101 (2017).

\bibitem{sophia}
Liu, Zhiyuan, et al. \\
``Sophia: A scalable stochastic second-order optimizer for language model pre-training.'' \\
arXiv preprint arXiv:2305.14342 (2023).
\end{thebibliography}

\end{document}
