\documentclass{article}

\usepackage{graphicx}
\usepackage{amsmath}

\title{Adaptive Orthogonal Momentum: A Novel Optimizer for Transformer Language Models}
\author{Aardvark}

\begin{document}

\maketitle

\begin{abstract}
We present Adaptive Orthogonal Momentum (AOM), a novel optimizer for transformer language models that combines selective orthogonalization with adaptive learning rates. AOM achieves a validation loss of 3.808 on the FineWeb benchmark, outperforming the AdamW baseline (4.927) and approaching the state-of-the-art Muon optimizer (3.537). Our key innovation is the integration of layer-specific orthogonal gradient processing with momentum-based adaptation, enabling more stable training and faster convergence. Extensive ablations demonstrate the effectiveness of our approach, particularly in attention layers where orthogonalization provides the most benefit.
\end{abstract}

\section{Introduction}
Optimizing transformer language models remains challenging due to the complex interaction between different architectural components and the need for stable, efficient training. While AdamW has become the de facto standard, recent work has shown that specialized optimizers can significantly improve performance. We introduce Adaptive Orthogonal Momentum (AOM), which addresses key limitations in existing approaches through:

\begin{itemize}
    \item Selective orthogonalization: Applying orthogonal gradient processing primarily to attention layers
    \item Layer-wise adaptation: Different learning rates and momentum for attention, FFN, and embedding layers
    \item Warmup scheduling: Gradual introduction of adaptive components for stability
\end{itemize}

\section{Related Work}
Our work builds on several key developments in optimization:

\subsection{Orthogonal Gradient Processing}
The Muon optimizer demonstrated the benefits of orthogonal gradient processing, particularly for attention layers. However, it applies orthogonalization uniformly across all parameters, which our results show to be suboptimal.

\subsection{Adaptive Optimization}
AdamW's success stems from its adaptive learning rates, but it lacks layer-specific adaptation. Recent work has explored per-layer learning rates, but without the stability benefits of orthogonalization.

\subsection{Momentum Methods}
Momentum has proven effective for deep learning optimization, but traditional approaches use a single momentum value across all parameters. Our layer-specific momentum adaptation provides better control over the optimization process.

\section{Method}
\subsection{Selective Orthogonalization}
We apply orthogonal gradient processing primarily to attention layers using a Newton-Schulz iteration with optimal coefficients. The number of orthogonal steps is adapted based on layer type:

\begin{equation}
    X_{k+1} = aX_k + (bA + cA^2)X_k
\end{equation}

where $A = X_kX_k^T$ and $a, b, c$ are optimized coefficients.

\subsection{Layer-wise Adaptation}
Different components use distinct learning rates and momentum values:

\begin{itemize}
    \item Attention layers: $\eta=0.01$, $\beta_1=0.9$
    \item FFN layers: $\eta=0.005$, $\beta_1=0.85$
    \item Embeddings: $\eta=0.001$, $\beta_1=0.9$
\end{itemize}

\subsection{Warmup Scheduling}
We gradually increase the strength of adaptive components over the first 1000 steps to prevent early instability.

\section{Experiments}
We evaluate AOM on the FineWeb benchmark using a Qwen 3 architecture with 134M parameters. Our baselines include AdamW and Muon. Training runs for 400 steps with a batch size of 256.

\section{Results}
\input{results}

\begin{figure}[h]
\centering
\includegraphics[width=0.8\textwidth]{figures/performance_comparison.pdf}
\caption{Comparison of final validation loss across different optimizers. Lower is better.}
\label{fig:comparison}
\end{figure}

Figure~\ref{fig:comparison} shows the final validation loss comparison between AOM and other optimizers. AOM achieves a final validation loss of 3.808, outperforming AdamW (4.927) and approaching Muon (3.537). Key findings:

\begin{itemize}
    \item Selective orthogonalization provides the most benefit in attention layers
    \item Layer-wise adaptation prevents over-constraining FFN layers
    \item Warmup scheduling significantly improves training stability
\end{itemize}

\section{Discussion}
While AOM shows promising results, there are several areas for improvement:

\begin{itemize}
    \item Memory overhead from orthogonalization could be reduced
    \item Dynamic adaptation of orthogonalization strength
    \item Better coordination between orthogonal and adaptive updates
\end{itemize}

\section{Conclusion}
We presented Adaptive Orthogonal Momentum, a novel optimizer that combines selective orthogonalization with layer-wise adaptation. Our results demonstrate significant improvements over AdamW while approaching state-of-the-art performance. Future work will focus on reducing memory overhead and improving dynamic adaptation.

\end{document}
