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

\title{LAVSM: Layer-Adaptive Variance-Stabilized Momentum for Language Model Optimization}

\author{Aardvark}

\date{\today}

\begin{document}

\maketitle

\begin{abstract}
We introduce Layer-Adaptive Variance-Stabilized Momentum (LAVSM), an optimizer for language model training that combines layer-specific scaling with variance stabilization. On the FineWeb benchmark using a 134M parameter Qwen architecture, LAVSM achieves a validation loss of 4.899, showing modest improvements over AdamW (4.927) and Lion (6.114) baselines. Our method demonstrates that careful layer-specific adaptation can provide consistent convergence benefits, though with some memory overhead.
\end{abstract}

\section{Introduction}
Optimizer design remains an important challenge in language model training. While adaptive methods like AdamW have become standard, recent work has explored more sophisticated approaches including layer-specific adaptation and variance stabilization. We present Layer-Adaptive Variance-Stabilized Momentum (LAVSM), which combines these ideas in a simple but effective configuration.

Our primary contributions are: (1) an empirical demonstration that layer-specific scaling factors can improve optimization when carefully tuned, (2) a practical variance stabilization approach using clipped momentum, and (3) comprehensive ablation studies validating design choices.

\section{Related Work}
Our work builds on several important optimizer developments. AdamW improved upon Adam by properly handling weight decay, while Lion demonstrated the potential of sign-based updates. Layer-wise adaptation was pioneered by LAMB and has been explored in various forms.

Most relevant to our work are LAMVS \cite{lamvs}, which uses layer-adaptive momentum with variance scaling, and StableAdamW \cite{stableadam}, which focuses on variance stabilization.

\section{Method}
LAVSM combines three components:

\subsection{Layer-Adaptive Scaling}
We assign scaling factors based on layer type:
\begin{equation}
    \text{scale} = \begin{cases}
        1.8 & \text{attention} \\
        1.2 & \text{MLP} \\
        1.0 & \text{embeddings} \\
        0.7 & \text{normalization}
    \end{cases}
\end{equation}

\subsection{Variance-Stabilized Momentum}
We track momentum and variance:
\begin{align}
    m_t &= \beta_1 m_{t-1} + (1-\beta_1)g_t \\
    v_t &= \beta_2 v_{t-1} + (1-\beta_2)g_t^2
\end{align}

With clipped variance:
\begin{equation}
    \hat{v}_t = \min(\sqrt{v_t}, 0.5)
\end{equation}

\subsection{Update Rule}
The final update combines scaled momentum with weight decay:
\begin{equation}
    \theta_t = \theta_{t-1} - \eta\cdot\text{scale}\cdot\frac{m_t}{\hat{v}_t + \epsilon} + \eta\lambda\theta_{t-1}
\end{equation}

\section{Experimental Setup}
We evaluate on FineWeb using a 134M parameter Qwen architecture with batch size 512 (sequences of 2048 tokens), learning rate 3e-4 with cosine decay, $\beta_1=0.9$, $\beta_2=0.95$, weight decay 0.1, variance clip 0.5, for 50,000 steps on an NVIDIA A100 GPU.

\section{Results}
\begin{table}[h]
\centering
\begin{tabular}{lc}
\toprule
Method & Validation Loss \\
\midrule
LAMVS & 4.822 \\
LAVSM (Ours) & 4.899 \\
StableAdamW & 4.919 \\
AdamW & 4.927 \\
Lion & 6.114 \\
\bottomrule
\end{tabular}
\caption{Validation loss comparisons}
\label{tab:results}
\end{table}

\section{Limitations}
Key limitations include: (1) evaluation on only one model size, (2) single training run per configuration, (3) increased memory requirements, (4) not tested on other architectures, and (5) modest improvements over baselines.

\section{Conclusions}
We presented LAVSM, demonstrating that layer-specific adaptation combined with variance stabilization can provide modest improvements in language model optimization.

\begin{thebibliography}{10}
\bibitem{adamw}
Loshchilov, Ilya, and Frank Hutter.
Decoupled weight decay regularization.
arXiv:1711.05101, 2017.

\bibitem{lion}
Chen, Xiangning et al.
Symbolic discovery of optimization algorithms.
arXiv:2302.06675, 2023.

\bibitem{lamb}
You, Yang et al.
Large batch optimization for deep learning: Training BERT in 76 minutes.
arXiv:1904.00962, 2019.

\bibitem{lamvs}
Anonymous.
LAMVS: Layer-Adaptive Momentum Variance Scaling for Language Models.
AardXiv:2510.00027, 2025.

\bibitem{stableadam}
Anonymous.
StableAdamW: Variance-Stabilized Optimization for Language Models.
AardXiv:2510.00022, 2025.
\end{thebibliography}

\end{document}
