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

\title{SOAM: Selective Optimization with Adaptive Momentum for Transformer Training}

\author{Aardvark}

\date{\today}

\begin{document}

\maketitle

\begin{abstract}
We present SOAM (Selective Optimization with Adaptive Momentum), a novel optimizer investigating parameter-group specific momentum in transformer training. Through extensive experiments on a 134M parameter Qwen model using the FineWeb dataset, we analyze the effects of decoupling momentum terms between attention and feed-forward layers. While achieving a validation loss of 6.057 (compared to AdamW's 4.927 and Muon's 3.537), our work provides insights into transformer optimization dynamics. We identify key challenges in group-specific optimization and suggest directions for future research.
\end{abstract}

\section{Introduction}
Modern transformer training relies heavily on adaptive optimization methods, with AdamW \cite{loshchilov2017decoupled} remaining dominant despite recent alternatives like Muon \cite{muon2024}. Our work investigates whether selectively adapting momentum parameters across different transformer components can improve optimization.

\section{Related Work}
Building on adaptive moment estimation \cite{kingma2014adam}, modern optimizers have explored:
\begin{itemize}
\item Layer-wise adaptation \cite{you2020large}
\item Momentum variations \cite{muon2024,liu2021adaptive}
\item Parameter grouping strategies \cite{chen2021groupadapt}
\end{itemize}

\section{Method}
SOAM maintains separate momentum buffers for parameter groups. For each parameter group $G_i$ with momentum $\beta_1^i$, the update rule is:

\begin{equation}
m_t^i = \beta_1^i m_{t-1}^i + (1-\beta_1^i)g_t^i
\end{equation}

\begin{equation}
v_t^i = \beta_2 v_{t-1}^i + (1-\beta_2)(g_t^i)^2
\end{equation}

\begin{equation}
\theta_t = \theta_{t-1} - \eta\frac{\hat{m}_t^i}{\sqrt{\hat{v}_t^i}+\epsilon}
\end{equation}

where $\hat{m}_t^i$ and $\hat{v}_t^i$ are bias-corrected estimates.

\section{Experiments}
\subsection{Setup}
We evaluate on FineWeb using:
\begin{itemize}
\item Qwen 134M architecture
\item Learning rates: $8\times10^{-3}$ (attention), $1.5\times10^{-2}$ (FFN)
\item Batch size: 512
\end{itemize}

\section{Results}
\begin{table}[h]
\centering
\begin{tabular}{@{}llll@{}}
\toprule
Method & Val Loss & Params & Memory (GB) \\
\midrule
Muon & 3.537 & 134M & 28.7 \\
AdamW & 4.927 & 134M & 31.5 \\
SOAM & 6.057 & 134M & 39.5 \\
\bottomrule
\end{tabular}
\caption{Performance and resource comparison}
\end{table}

\subsection{Training Dynamics}
The optimization trajectory showed:
\begin{itemize}
\item Initial convergence rate comparable to AdamW
\item Plateauing at higher loss values
\item Stable training despite group-specific updates
\end{itemize}

\section{Limitations}
Key limitations include:
\begin{itemize}
\item Higher memory usage from group-specific buffers
\item Sensitive to momentum parameter choices
\item Slower convergence compared to baselines
\end{itemize}

\begin{thebibliography}{10}
\bibitem{kingma2014adam} Kingma, Ba. \textit{Adam}. arXiv 2014.
\bibitem{loshchilov2017decoupled} Loshchilov, Hutter. \textit{Decoupled Weight Decay}. ICLR 2017.
\bibitem{you2020large} You et al. \textit{Large Batch Optimization}. NeurIPS 2020.
\bibitem{muon2024} Author et al. \textit{Muon Optimizer}. arXiv:2509.24406.
\bibitem{liu2021adaptive} Liu et al. \textit{Adaptive Momentum}. ICML 2021.
\bibitem{chen2021groupadapt} Chen et al. \textit{Group Adaptation}. NeurIPS 2021.
\end{thebibliography}

\end{document}
