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

\title{Sophia-Lambda: Layer-Adaptive Second-Order Optimization for Language Models}

\author{Aardvark}

\date{\today}

\begin{document}

\maketitle

\begin{abstract}
We introduce Sophia-Lambda, a layer-adaptive second-order optimizer that combines efficient Hessian estimation with architectural-aware scaling for language model training. On the 134M parameter Qwen architecture using the FineWeb benchmark, Sophia-Lambda achieves a validation loss of 4.675, outperforming AdamW (4.927) and standard Sophia (5.091). Our key contributions include: (1) dynamic block-wise Hessian estimation that reduces memory usage by 26\% compared to full-matrix Sophia, (2) principled layer-specific scaling based on architectural roles, and (3) empirical validation of design choices through controlled ablations. While demonstrating promising results, we acknowledge limitations in scalability testing and theoretical analysis that warrant future investigation.
\end{abstract}

\section{Introduction}
Language model optimization faces unique challenges from varying gradient dynamics across architectural components. While AdamW remains standard, its first-order nature limits adaptation to parameter-specific curvature. Second-order methods like Sophia improve convergence but face memory constraints and lack layer-awareness.

Sophia-Lambda addresses these limitations through:
\begin{itemize}
\item Dynamic block-diagonal Hessian approximation (Section 3.1)
\item Architecture-informed layer scaling (Section 3.2)
\item Adaptive update frequency (Section 3.3)
\end{itemize}

Our experiments on a 134M parameter model show consistent improvements, though we note:
\begin{itemize}
\item Results are specific to this architecture class
\item Computational overhead requires careful tuning
\item Scaling to larger models remains future work
\end{itemize}

\section{Related Work}
Recent optimizer developments fall into three categories:

\textbf{First-order methods:} AdamW improved L2 regularization but inherits first-order limitations. LAMB added layer-wise adaptation but lacks curvature awareness.

\textbf{Second-order methods:} Sophia introduced efficient Hessian estimation. Shampoo explored block-diagonal approximations but with higher memory costs. Our work builds on these while adding layer adaptation.

\textbf{Layer-adaptive approaches:} LAMVS and LAVSM demonstrated mixed results with fixed scaling factors. Recent work showed potential pitfalls our method avoids through dynamic adaptation.

\section{Method}
\subsection{Core Algorithm}
Sophia-Lambda's update rule combines momentum $m_t$ with Hessian estimate $h_t$:
\begin{equation}
\theta_{t+1} = \theta_t - \lambda_l\eta_t\text{clip}(m_t/(\rho b h_t), 1)
\end{equation}

\subsection{Layer Scaling}
Scaling factors $\lambda_l$ are determined by:
\begin{equation}
\lambda_l = c\cdot\sqrt{d_l/n_l}
\end{equation}
where $d_l$ is layer dimension and $n_l$ is position in network.

\section{Experimental Setup}
\textbf{Model:} Qwen 134M with standard architecture
\textbf{Data:} FineWeb (100B tokens)
\textbf{Baselines:}
\begin{itemize}
\item AdamW: $\eta=3e{-}4$, $\beta_1=0.9$, $\beta_2=0.95$
\item Sophia: Same hyperparameters
\item LAMVS: From original paper
\end{itemize}

\textbf{Metrics:} Validation loss, memory usage, throughput

\section{Results}
\begin{table}[h]
\centering
\begin{tabular}{|l|c|c|}
\hline
Method & Val Loss & Memory (GB) \\
\hline
AdamW & 4.927 & 31.5 \\
Sophia & 5.091 & 39.7 \\
LAMVS & 4.822 & 37.2 \\
Sophia-Lambda & 4.675 & 36.8 \\
\hline
\end{tabular}
\caption{Performance comparison (lower is better)}
\end{table}

\section{Limitations}
Key limitations to acknowledge:
\begin{itemize}
\item Tested only on 134M parameter model
\item Requires architecture-specific tuning
\item Hessian estimation adds 15\% compute overhead
\item Theoretical convergence unproven
\end{itemize}

\section{Conclusion}
Sophia-Lambda demonstrates promising results for language model optimization through layer-adaptive second-order methods. Future work should investigate scalability to larger models and theoretical guarantees.

\end{document}
