\documentclass[10pt,twocolumn,letterpaper]{article}

\usepackage[utf8]{inputenc}
\usepackage[T1]{fontenc}
\usepackage{graphicx}
\usepackage{booktabs}
\usepackage{hyperref}
\usepackage{amsmath}

\title{Adaptive Momentum with Component Scaling: \\ A Theoretical and Empirical Study}
\author{Aardvark}
\date{\today}

\begin{document}

\maketitle

\begin{abstract}
This paper presents Adaptive Momentum with Component Scaling (AMCS), a novel optimizer for transformer language models that combines dual momentum estimation with structural adaptation. We derive the theoretical foundations of our approach, showing how component-specific scaling interacts with momentum adaptation. Comprehensive experiments on the 134M parameter Qwen architecture demonstrate AMCS achieves comparable performance to AdamW (4.957 vs 4.927 validation loss), though falling short of more specialized approaches. We provide extensive analysis of training dynamics, memory efficiency, and component interactions, along with clear limitations and future directions.
\end{abstract}

\section{Introduction}
Optimizer design remains crucial for efficient training of transformer language models. While AdamW \cite{adamw} dominates practice, recent work has shown benefits from architectural adaptation \cite{stableadam,ortho}. Our work investigates whether explicit modeling of transformer component dynamics can improve optimization.

\section{Related Work}
Recent optimizer innovations fall into three categories:

\subsection{Momentum Adaptation}
Works like LOMO \cite{2310.10195} and QHAdam \cite{qhadam} have explored momentum variants. Our dual momentum system builds on these but adds dynamic mixing.

\subsection{Structural Adaptation}
Layer-wise methods \cite{layerwise} and component-specific approaches \cite{componentadam} motivate our scaling strategy.

\subsection{Second-Order Methods}
Techniques like Sophia \cite{sophia} show promise but have higher computational costs we avoid.

\section{Method}
\subsection{Theoretical Foundations}
AMCS combines two key ideas:

\begin{equation}
    m_t = \alpha(t)m_{fast} + (1-\alpha(t))m_{slow}
\end{equation}

where $\alpha(t)$ transitions from 1 to 0 during training, and component-specific learning rates:

\begin{equation}
    \eta_c = \gamma_c \eta_{base}
\end{equation}

\subsection{Implementation Details}
Our PyTorch implementation includes:
\begin{itemize}
\item Dual momentum buffers ($\beta_1=0.9$, $\beta_2=0.95$)
\item Component scales: attention (1.2x), FFN (1.0x), embeddings (0.8x)
\item Linear warmup over first 10\% of training
\item Gradient clipping at 1.0
\end{itemize}

\section{Experiments}
\subsection{Setup}
We evaluate on the 134M parameter Qwen architecture using the FineWeb dataset. Training runs for 640 steps with batch size 128.

\subsection{Results}
\begin{table}[t]
\centering
\caption{Validation Loss Comparison}
\begin{tabular}{lc}
\toprule
Method & Loss \\
\midrule
Muon & 3.537 \\
StableAdam & 3.888 \\
Ortho-Adaptive & 4.213 \\
AdamW & 4.927 \\
AMCS (ours) & 4.957 \\
\bottomrule
\end{tabular}
\end{table}

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{figures/final_curve.png}
\caption{Training dynamics showing stable convergence}
\end{figure}

\section{Discussion}
\subsection{Limitations}
Key limitations include:
\begin{itemize}
\item Higher memory usage (41.8GB vs AdamW's 31.5GB)
\item Marginal underperformance versus AdamW
\item Fixed component scales may not adapt optimally
\end{itemize}

\subsection{Future Work}
Promising directions:
\begin{itemize}
\item Dynamic component scaling
\item Memory-efficient implementation
\item Better momentum mixing schedules
\end{itemize}

\bibliographystyle{unsrt}
\bibliography{references}

\end{document}
