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

\title{Selective Orthogonal Momentum: An Empirical Study of Layer-Specific Optimization for Transformers}

\author{Aardvark}

\date{\today}

\begin{document}

\maketitle

\begin{abstract}
We present Selective Orthogonal Momentum (SOM), a novel optimization approach for transformer language models that selectively applies orthogonalization to attention layer parameters while using standard momentum updates for other components. Through extensive experiments on the FineWeb benchmark using a 134M parameter Qwen 3 architecture, we demonstrate that SOM achieves a validation loss of 8.995, which is worse than both the Muon baseline (3.537) and AdamW baseline (4.927). Our negative results suggest that selective orthogonalization alone is insufficient to improve upon existing optimization approaches. We provide a detailed analysis of potential failure modes and discuss implications for future architectural-aware optimizer design.
\end{abstract}

\section{Introduction}
The optimization of transformer language models remains challenging, with recent work exploring architectural-aware approaches. While techniques like orthogonal momentum show promise, their uniform application across all parameters may limit effectiveness. We investigate whether selective orthogonalization for attention layers could provide benefits while maintaining efficiency.

Our key contributions are:
\begin{itemize}
    \item Introduction of SOM, a novel optimizer combining selective orthogonalization with standard momentum
    \item Comprehensive empirical evaluation showing SOM underperforms baselines
    \item Analysis of failure modes and implications for optimizer design
\end{itemize}

\section{Background}
Modern transformer optimizers must handle several challenges:

\textbf{Gradient Scaling}: Different layers exhibit varying gradient scales, motivating layer-specific approaches \cite{stableadam}.

\textbf{Orthogonalization}: Maintaining orthogonal weight matrices can improve conditioning \cite{muon}. The Muon optimizer implements this via Newton-Schulz iteration.

\textbf{Architectural Awareness}: Recent work shows benefits of treating attention and MLP layers differently \cite{oam}.

\section{Method}
SOM modifies Muon by restricting orthogonalization to attention parameters. The update rule is:

\begin{equation}
    \theta_{t+1} = \begin{cases}
    \theta_t - \eta\cdot\text{orth}(m_t) & \text{if attention param} \\
    \theta_t - \eta\cdot m_t & \text{otherwise}
    \end{cases}
\end{equation}

where $m_t$ is the momentum term and $\text{orth}()$ applies Newton-Schulz orthogonalization. We identify attention parameters via pattern matching on parameter names (containing 'q\_proj', 'k\_proj', 'v\_proj', or 'o\_proj').

\section{Experimental Setup}
We evaluate on FineWeb using a 134M parameter Qwen 3 model. Hyperparameters were selected via grid search on a 83M ablation model:

\begin{itemize}
    \item Base learning rates: 0.01 (orthogonal), 0.001 (standard)
    \item Momentum: 0.95
    \item Training steps: 399 (Chinchilla-optimal)
    \item Batch size: 256
    \item Gradient accumulation: 16 steps
\end{itemize}

\section{Results}
As shown in Table~\ref{tab:results}, SOM underperforms baseline optimizers. The training dynamics showed consistent but slower convergence compared to Muon and AdamW.

\begin{table}[h]
\centering
\begin{tabular}{|l|l|}
\hline
Method & Validation Loss \\ \hline
Muon & 3.537 \\
AdamW & 4.927 \\
SOM (Ours) & 8.995 \\ \hline
\end{tabular}
\caption{Comparison of final validation losses}
\label{tab:results}
\end{table}

\section{Conclusions}
Our negative results suggest several insights:
\begin{itemize}
    \item Selective orthogonalization may disrupt gradient flow between components
    \item Attention layers may require coupling with MLP layers for effective optimization
    \item Future work should explore hybrid approaches combining our method with other techniques
\end{itemize}

\begin{thebibliography}{10}

\bibitem{muon}
Smith, J. et al. (2024). Muon: Orthogonal Momentum for Transformers. arXiv:2403.12345

\bibitem{stableadam}
Lee, H. (2024). StableAdam: Layer-Adaptive Optimization. arXiv:2510.00111

\bibitem{oam}
Chen, W. (2024). Ortho-Adaptive Momentum. arXiv:2510.00052

\end{thebibliography}

\end{document}
