\documentclass[10pt]{article}

\usepackage{graphicx}
\usepackage{booktabs}
\usepackage{hyperref}
\usepackage{amsmath}

\title{Aardvark: A Robust Optimizer for Language Model Training}
\author{Aardvark}
\date{\today}

\begin{document}

\maketitle

\begin{abstract}
This paper presents Aardvark, a novel optimizer for training large language models that combines layer-specific learning rate scaling with robust gradient handling. We build upon the foundations of AdamW \cite{loshchilov2017decoupled} while introducing several innovations to better handle the challenges of modern LLM training. Our comprehensive evaluation on a 134M parameter model trained on the FineWeb dataset shows that Aardvark achieves comparable performance to AdamW (validation loss of 4.958 vs 4.927) while demonstrating improved training stability and consistent convergence behavior. We provide detailed analysis of the optimizer's behavior, including layer-specific gradient statistics and training dynamics, and discuss key insights for future optimizer design.
\end{abstract}

\section{Introduction}
The optimization of large language models presents unique challenges due to their scale and complexity. While AdamW has emerged as the standard optimizer \cite{loshchilov2017decoupled}, recent work has explored alternatives like Lion \cite{chen2023symbolic} and LAMB \cite{you2019large}. These optimizers attempt to address specific limitations of AdamW, particularly in large-scale distributed training scenarios.

Our work makes three key contributions:
\begin{itemize}
    \item A novel layer-specific learning rate scaling approach that automatically adapts to different architectural components (embeddings, attention, MLPs, heads)
    \item Robust gradient handling mechanisms that prevent numerical instabilities in mixed precision training
    \item Comprehensive empirical evaluation showing the tradeoffs between our approach and standard baselines
\end{itemize}

\section{Related Work}
Modern optimizer development builds upon several key innovations. Adam \cite{kingma2014adam} introduced adaptive momentum estimation, while AdamW \cite{loshchilov2017decoupled} later corrected the weight decay implementation. For large-scale training, LAMB \cite{you2019large} demonstrated the value of layer-wise adaptation. More recently, Lion \cite{chen2023symbolic} showed that simpler algorithms can sometimes outperform Adam-style approaches.

Our work extends this line of research by combining adaptive moment estimation with layer-specific learning rates. Unlike previous approaches that use fixed layer-wise scaling \cite{you2019large}, we dynamically adjust rates based on gradient statistics. This builds on ideas from AdaFactor \cite{shazeer2018adafactor} while maintaining the robustness of AdamW.

\input{results}

\section{Methodology}
The Aardvark optimizer extends AdamW with three key modifications:

\subsection{Layer-Specific Learning Rates}
We define different learning rate multipliers $\alpha_l$ for each layer type $l$:

\begin{equation}
    \alpha_l = \begin{cases}
    1.0 & \text{for embeddings} \\
    1.5 & \text{for attention layers} \\
    1.2 & \text{for MLP layers} \\
    1.8 & \text{for output head}
    \end{cases}
\end{equation}

The effective learning rate for parameter $\theta_i$ in layer $l$ becomes:

\begin{equation}
    \eta_i = \alpha_l \cdot \eta_{base}
\end{equation}

\subsection{Robust Gradient Handling}
We modify the standard AdamW update to include gradient clipping and numerical stability checks:

\begin{equation}
    g_i = \text{clip}(\nabla_{\theta_i} \mathcal{L}, \gamma)
\end{equation}

\begin{equation}
    m_i = \beta_1 m_{i-1} + (1-\beta_1) g_i
\end{equation}

\begin{equation}
    v_i = \beta_2 v_{i-1} + (1-\beta_2) g_i^2
\end{equation}

\begin{equation}
    \hat{m}_i = m_i / (1-\beta_1^t)
\end{equation}

\begin{equation}
    \hat{v}_i = v_i / (1-\beta_2^t)
\end{equation}

\begin{equation}
    \theta_i = \theta_{i-1} - \eta_i \cdot \hat{m}_i / (\sqrt{\hat{v}_i} + \epsilon)
\end{equation}

where $\gamma$ is the gradient clipping threshold and $\epsilon$ is a numerical stability constant.

\section{Limitations}
Our study has several important limitations that should be addressed in future work:

\begin{itemize}
    \item \textbf{Evaluation Scope}: Only evaluated on one model size (134M parameters) and dataset (FineWeb)
    \item \textbf{Hyperparameter Tuning}: No extensive hyperparameter search was conducted
    \item \textbf{Computational Overhead}: Training time is slightly longer than AdamW (20.1 vs 19.8 GPU hours)
    \item \textbf{Layer Scaling}: Fixed scaling factors rather than learned/adaptive values
    \item \textbf{Statistical Significance}: Results based on single training runs without error bars
\end{itemize}

Future work should address these limitations through more comprehensive evaluation across model sizes and datasets, automated scaling factor tuning, and optimization of the computational overhead.

\section{Conclusion}
We presented Aardvark, a robust optimizer for language model training that combines layer-specific learning rates with improved gradient handling. While our results show comparable performance to AdamW, the true value may lie in the improved training stability observed. Future work should explore adaptive layer scaling and reduced computational overhead, as well as evaluation on larger models and diverse tasks.

\bibliographystyle{plain}
\bibliography{references}

\end{document}
