\documentclass{article}
\usepackage{graphicx}
\usepackage{amsmath}

\title{HyMo: A Study of Hybrid Momentum Optimization for Transformer Language Models}
\author{Aardvark}
\date{\today}

\begin{document}

\maketitle

\begin{abstract}
This paper presents a systematic investigation of hybrid momentum optimization techniques for transformer language models. We examine the feasibility of combining standard momentum updates with selective orthogonalization for large parameter matrices, focusing on training stability and performance tradeoffs. Our experiments on a 134M parameter transformer model demonstrate that while our HyMo optimizer achieves comparable performance to AdamW (validation loss of 4.983 vs 4.927), it does not outperform existing approaches. The study provides insights into the practical challenges of incorporating orthogonal updates in modern language model training pipelines and establishes baseline expectations for similar hybrid approaches.
\end{abstract}

\section{Introduction}
Optimizer design remains an active area of research in deep learning, particularly for large language models where training stability and efficiency are paramount. While Adam and its variants dominate current practice~\cite{kingma2014adam,loshchilov2017decoupled}, recent work has explored alternative approaches including orthogonal optimization methods~\cite{liu2023muon} and second-order techniques~\cite{liu2023sophia}.

This work investigates whether selective orthogonalization can be beneficially combined with standard momentum updates in a hybrid approach. Our primary research questions are:
\begin{itemize}
\item Can selective orthogonalization maintain training stability while potentially offering benefits over pure momentum methods?
\item What are the practical tradeoffs in computational overhead and performance when implementing such hybrid approaches?
\item How does selective orthogonalization compare to full orthogonalization methods like Muon?
\end{itemize}

\section{Related Work}
Our work builds on several established lines of research in optimization for deep learning. The success of Adam~\cite{kingma2014adam} demonstrated the effectiveness of adaptive momentum methods with per-parameter learning rates. Recent variants like AdamW~\cite{loshchilov2017decoupled} introduced improved weight decay handling.

Orthogonal optimization approaches have shown promise in various contexts~\cite{wisdom2016full,liu2023muon}. The Muon optimizer demonstrated that orthogonal momentum could improve training efficiency, though with scalability challenges. Second-order methods like Sophia~\cite{liu2023sophia} have shown potential but require careful implementation to remain practical.

Our work differs by focusing on a targeted combination of techniques rather than proposing a fundamentally new optimization approach. This aligns with recent work on optimizer composition~\cite{chen2023composer} while maintaining computational efficiency.

\section{Method}
The HyMo optimizer combines standard momentum updates with selective orthogonalization for large parameter matrices. The complete procedure is:

1. Compute gradient $g_t \leftarrow \nabla_\theta L(\theta_{t-1})$
2. Update momentum: $m_t \leftarrow \beta_1 m_{t-1} + (1-\beta_1)g_t$
3. Update variance: $v_t \leftarrow \beta_2 v_{t-1} + (1-\beta_2)g_t^2$
4. For large matrices (dim $\geq1024\times1024$):
   a. Orthogonalize: $m_t^{orth} \leftarrow 0.5(m_t - m_t m_t^T m_t)$
   b. Update: $\theta_t \leftarrow \theta_{t-1} - \eta m_t^{orth}/\sqrt{v_t + \epsilon}$
5. For other parameters:
   a. Update: $\theta_t \leftarrow \theta_{t-1} - \eta m_t/\sqrt{v_t + \epsilon}$

The 1024 threshold was chosen empirically based on preliminary experiments showing diminishing returns for orthogonalization on smaller matrices. The conservative orthogonalization step helps maintain stability while potentially benefiting from orthogonal updates.

\section{Experiments}
We evaluated HyMo on a 134M parameter Qwen transformer model trained on the FineWeb dataset. All experiments used:
\begin{itemize}
\item Batch size: 4M tokens
\item Learning rate: 3e-4 (with linear warmup over first 100 steps)
\item Weight decay: 0.1
\item $\beta_1,\beta_2$: 0.9, 0.98
\item Training steps: 640
\item Hardware: 8x A100 GPUs with FSDP
\end{itemize}

We compared against AdamW and Muon baselines using identical training configurations. The Muon implementation followed~\cite{liu2023muon} with default hyperparameters.

\section{Results}
\begin{figure}[h]
\centering
\includegraphics[width=0.8\linewidth]{figures/training_curve.png}
\caption{Training loss curves showing HyMo's convergence behavior}
\label{fig:training}
\end{figure}

Figure~\ref{fig:training} shows HyMo's training curve, demonstrating stable convergence. The final validation losses were:

\begin{center}
\begin{tabular}{|l|c|}
\hline
Method & Validation Loss \\
\hline
Muon & 3.537 \\
AdamW & 4.927 \\
HyMo (Ours) & 4.983 \\
\hline
\end{tabular}
\end{center}

Key observations:
\begin{itemize}
\item HyMo maintains comparable performance to AdamW
\item The orthogonalization overhead (41.78GB memory vs AdamW's 31.49GB) may not justify the minimal performance difference
\item Muon's superior performance suggests full orthogonalization may be preferable when computationally feasible
\end{itemize}

\section{Discussion}
Our results suggest that selective orthogonalization, while feasible, does not provide clear advantages over either pure momentum methods or full orthogonalization approaches. Several factors may contribute:

1. The benefits of orthogonalization may require more aggressive application than our conservative approach
2. The computational overhead may outweigh any theoretical benefits
3. Modern architectures may be sufficiently robust to not require orthogonal updates

These findings align with recent work suggesting that careful initialization and normalization may reduce the need for specialized optimizers~\cite{zhang2023stabilizing}.

\section{Conclusion}
This study investigated hybrid momentum optimization for transformer language models. While our HyMo approach demonstrated feasibility, it did not outperform existing methods. The work provides valuable negative results and establishes baseline expectations for similar hybrid approaches. Future work could explore more sophisticated criteria for applying orthogonal updates or investigate whether the benefits of orthogonalization can be achieved through other means.

\bibliographystyle{unsrt}
\bibliography{references}

\end{document}
