papers/zen/zen-aso-protocol.tex

\documentclass[11pt,a4paper]{article}
\usepackage[utf8]{inputenc}
\usepackage[T1]{fontenc}
\usepackage{amsmath,amsfonts,amssymb}
\usepackage{amsthm}
\newtheorem{theorem}{Theorem}
\newtheorem{lemma}[theorem]{Lemma}
\newtheorem{definition}[theorem]{Definition}
\usepackage{graphicx}
\usepackage{hyperref}
\usepackage{listings}
\usepackage{color}
\usepackage{booktabs}
\usepackage{float}
\usepackage{geometry}
\usepackage{algorithm}
\usepackage{algorithmic}
\geometry{margin=1in}
\definecolor{zenblue}{RGB}{41,121,255}
\hypersetup{colorlinks=true,linkcolor=zenblue,urlcolor=zenblue,citecolor=zenblue}

\title{\textbf{ASO: Active Semantic Optimization for Distributed AI}\\
\large Technical Report v2025.05}
\author{Antje Worring, Hanzo AI Research, Zen LM Research Team\\
\texttt{[email protected]}}
\date{May 2025}

\begin{document}
\maketitle

\begin{abstract}
We introduce Active Semantic Optimization (ASO), a distributed training protocol defined
under Hanzo Improvement Proposal HIP-002. ASO replaces conventional token-level gradient
aggregation with a semantics-aware update scheme in which gradient steps are weighted by
embedding-space proximity to a set of dynamically maintained \emph{semantic anchor points}.
This yields faster convergence on knowledge-intensive tasks, measurably reduces catastrophic
forgetting during continual learning, and integrates naturally with the Zen MoDE
(Mixture of Distilled Experts) architecture. We provide formal convergence guarantees under
standard smoothness assumptions, an empirical study across four distributed training regimes,
and an open reference implementation compatible with any transformer-based model.
\end{abstract}

\tableofcontents
\newpage

%% -----------------------------------------------------------------------
\section{Introduction}
\label{sec:intro}
%% -----------------------------------------------------------------------

Large-scale language model training distributes gradient computation across hundreds to
thousands of accelerators. Standard data-parallel and model-parallel schemes aggregate
raw parameter gradients, treating every training token as equally informative regardless
of its semantic content. This ignores a fundamental asymmetry in natural language: tokens
that carry high semantic load (named entities, numeric facts, causal connectives) contribute
disproportionately to downstream task performance, while high-frequency function words
contribute comparatively little per-gradient-step.

Active Semantic Optimization (ASO, HIP-002) addresses this asymmetry through three
coordinated mechanisms:

\begin{enumerate}
    \item \textbf{Semantic Loss Surfaces}: a reformulation of the training objective that
          modulates per-sample loss by a semantic relevance score derived from embedding
          proximity to a maintained anchor set.
    \item \textbf{Gradient Alignment via Embedding Proximity}: a projection step that
          suppresses gradient components that are orthogonal to the current semantic
          manifold, concentrating learning signal on semantically meaningful directions.
    \item \textbf{Adaptive Anchor Maintenance}: an online algorithm that tracks the
          evolving distribution of high-value semantics across the training corpus,
          ensuring anchors remain informative throughout training.
\end{enumerate}

ASO is designed for the Zen MoDE architecture \cite{zenlm2025mode} but applies to any
model with a dense embedding space. The protocol is implemented as a thin wrapper around
standard distributed training frameworks (PyTorch DDP, DeepSpeed ZeRO) and introduces
less than 3\% computational overhead relative to baseline training.

\subsection{Motivation}

Consider a factual recall task. At any given training step, the model receives a batch
containing tokens of wildly varying informational density: punctuation, stop-words,
proper nouns, numeric values, and abstract predicates. A uniform gradient aggregation
rule allocates equal credit to each token's contribution to the parameter update. This
is provably suboptimal when the downstream evaluation metric correlates strongly with
a non-uniform subset of training signal.

The ASO hypothesis, validated empirically in Section~\ref{sec:experiments}, is:

\begin{quote}
\emph{Routing gradient mass toward semantically dense regions of the embedding manifold
accelerates convergence on knowledge-intensive benchmarks without degrading fluency or
general language modeling perplexity.}
\end{quote}

%% -----------------------------------------------------------------------
\section{Background and Related Work}
\label{sec:background}
%% -----------------------------------------------------------------------

\subsection{Distributed Gradient Aggregation}

Standard distributed training minimizes the empirical risk:
\begin{equation}
    \mathcal{L}(\theta) = \frac{1}{N} \sum_{i=1}^{N} \ell(f_\theta(x_i), y_i)
    \label{eq:empirical_risk}
\end{equation}
using synchronous SGD across $K$ workers, where each worker $k$ computes a local gradient
$g_k = \nabla_\theta \mathcal{L}_k(\theta)$ and the server aggregates via
$g = \frac{1}{K}\sum_k g_k$.

Prior work on importance weighting (e.g., curriculum learning \cite{bengio2009curriculum},
self-paced learning \cite{kumar2010selfpaced}) re-weights samples at the loss level.
ASO extends this to the gradient level, applying weights that are \emph{content-adaptive}
and \emph{embedding-space grounded}.

\subsection{Semantic Representations in Training}

Embedding-based curriculum methods \cite{hacohen2019power} use fixed pretrained encoders
to order training examples by difficulty. ASO differs in two key respects: (1) the
semantic scorer is jointly updated with the main model parameters, and (2) the weighting
operates at gradient projection time rather than at batch construction time, enabling
fine-grained per-token control.

\subsection{Mixture of Experts and Routing}

The Zen MoDE architecture employs a sparse mixture-of-experts (MoE) routing layer that
assigns tokens to specialized sub-networks. ASO complements MoE routing by providing a
global semantic signal that can inform both the routing policy and the gradient update
for each expert, reducing expert collapse during early training.

%% -----------------------------------------------------------------------
\section{Active Semantic Optimization: Formulation}
\label{sec:formulation}
%% -----------------------------------------------------------------------

\subsection{Semantic Anchor Set}

Let $\mathcal{A} = \{a_1, \ldots, a_M\} \subset \mathbb{R}^d$ denote a set of $M$
semantic anchor vectors maintained in the model's embedding space $\mathbb{R}^d$.
Anchors are initialized from a \emph{semantic seed corpus} — a curated collection of
high-value factual and reasoning examples — and updated online via an exponential moving
average:
\begin{equation}
    a_j \leftarrow (1-\alpha) a_j + \alpha \cdot e(x^*)
    \label{eq:anchor_update}
\end{equation}
where $e(x^*)$ is the contextual embedding of the most semantically similar token in
the current batch to anchor $a_j$, and $\alpha \in (0,1)$ is the anchor learning rate.

\subsection{Semantic Relevance Score}

For a token $x_i$ with embedding $e_i \in \mathbb{R}^d$, define its semantic relevance
score as the soft-max proximity to the anchor set:
\begin{equation}
    s_i = \frac{1}{M} \sum_{j=1}^{M}
          \text{softmax}\!\left(\frac{e_i \cdot a_j}{\|e_i\|\|a_j\|\,\tau}\right)
    \label{eq:relevance_score}
\end{equation}
where $\tau > 0$ is a temperature parameter. This score lies in $(0,1)$ and approaches
1 when $e_i$ is closely aligned with some anchor and 0 when it is equidistant from all
anchors.

\subsection{Semantic Loss Surface}

The ASO training objective modulates the standard cross-entropy loss by the semantic
relevance scores:
\begin{equation}
    \mathcal{L}_{\text{ASO}}(\theta) =
        \frac{1}{N}\sum_{i=1}^{N}
        \bigl(1 + \lambda \cdot s_i\bigr)\,\ell\bigl(f_\theta(x_i), y_i\bigr)
    \label{eq:aso_loss}
\end{equation}
where $\lambda \geq 0$ controls the semantic emphasis. When $\lambda = 0$,
$\mathcal{L}_{\text{ASO}}$ reduces to the standard empirical risk.

\subsection{Gradient Alignment via Embedding Proximity}

Let $g_i = \nabla_\theta \ell(f_\theta(x_i), y_i)$ be the per-token gradient.
Define the \emph{semantic subspace} $\mathcal{S}$ as the span of the top-$R$ principal
components of $\{e_i \cdot a_j : j = 1,\ldots,M\}$, projected back into parameter space
via the Jacobian of the embedding layer. The ASO gradient is:
\begin{equation}
    \tilde{g}_i = s_i \cdot \text{Proj}_{\mathcal{S}}(g_i) +
                  (1 - s_i) \cdot g_i
    \label{eq:projected_gradient}
\end{equation}

Tokens with high semantic relevance ($s_i \approx 1$) have their gradients projected
onto the semantic subspace, concentrating the update in directions that improve
semantic fidelity. Tokens with low relevance ($s_i \approx 0$) pass through
unmodified, preserving general language modeling capability.

\subsection{Distributed Aggregation under ASO}

In a $K$-worker distributed setting, each worker $k$ computes local ASO gradients
$\{\tilde{g}_i\}$ for its mini-batch and reports the weighted aggregate:
\begin{equation}
    G_k = \frac{1}{|B_k|} \sum_{i \in B_k} \tilde{g}_i
    \label{eq:local_aggregate}
\end{equation}

The parameter server combines worker gradients using the semantics-weighted mean:
\begin{equation}
    G = \frac{\sum_k w_k G_k}{\sum_k w_k}, \quad
    w_k = \frac{1}{|B_k|} \sum_{i \in B_k} s_i
    \label{eq:global_aggregate}
\end{equation}

Workers whose batches contain more semantically dense content are given proportionally
higher weight in the global update, creating a soft curriculum at the batch level.

%% -----------------------------------------------------------------------
\section{Convergence Analysis}
\label{sec:theory}
%% -----------------------------------------------------------------------

\subsection{Assumptions}

We analyze ASO under standard non-convex optimization assumptions:

\begin{itemize}
    \item[\textbf{A1.}] $\mathcal{L}_{\text{ASO}}$ is $L$-smooth: $\|\nabla \mathcal{L}_{\text{ASO}}(\theta) - \nabla \mathcal{L}_{\text{ASO}}(\theta')\| \leq L\|\theta - \theta'\|$.
    \item[\textbf{A2.}] Stochastic gradient variance is bounded: $\mathbb{E}\|\tilde{g} - \nabla \mathcal{L}_{\text{ASO}}\|^2 \leq \sigma^2$.
    \item[\textbf{A3.}] The anchor set is updated with step size $\alpha \leq \frac{1}{2L}$.
    \item[\textbf{A4.}] The semantic subspace projection is $\rho$-approximately correct: the projection error satisfies $\|\text{Proj}_{\mathcal{S}}(g) - g_{\mathcal{S}}\| \leq \rho \|g\|$ for all $g$.
\end{itemize}

\subsection{Main Convergence Theorem}

\begin{theorem}[ASO Convergence Rate]
\label{thm:convergence}
Under Assumptions A1--A4, with learning rate $\eta = \sqrt{\frac{1}{LT}}$, after $T$
gradient steps the ASO iterate $\theta_T$ satisfies:
\begin{equation}
    \frac{1}{T}\sum_{t=1}^{T} \mathbb{E}\|\nabla \mathcal{L}_{\text{ASO}}(\theta_t)\|^2
    \leq
    \frac{2\sqrt{L(\mathcal{L}_{\text{ASO}}(\theta_0) - \mathcal{L}^*)}}{\sqrt{T}}
    + \sigma\sqrt{\frac{L}{T}} + \frac{\lambda \rho \sigma^2}{L}
    \label{eq:convergence}
\end{equation}
where $\mathcal{L}^*$ is the global minimum and the final term arises from projection
approximation error, vanishing as $\rho \to 0$.
\end{theorem}

\begin{proof}[Proof Sketch]
The proof follows the standard non-convex SGD analysis with two modifications: (1) the
semantic weighting introduces a factor of $(1 + \lambda s_i)$ in the descent lemma,
which is absorbed into the smoothness constant, and (2) the projection step incurs an
additive error bounded by $\lambda \rho \sigma^2 / L$ via a triangle inequality argument
on the projected gradient variance. Full proof in Appendix~A.
\end{proof}

\subsection{Implication: Effective Sample Complexity}

The effective sample complexity to reach $\epsilon$-stationarity under ASO is
$O(1/\epsilon^2)$, matching standard SGD. However, the constant factor is reduced by
the semantic emphasis term $\lambda$, which amplifies the contribution of high-signal
samples, yielding a practical speedup on semantically structured benchmarks.

%% -----------------------------------------------------------------------
\section{Implementation}
\label{sec:implementation}
%% -----------------------------------------------------------------------

\subsection{Anchor Initialization}

Anchors are initialized by running a forward pass over the semantic seed corpus
(typically 1--5M tokens of curated factual text) and clustering the resulting embeddings
with $k$-means ($k = M = 512$ in all experiments). Initialization takes approximately
15 minutes on a single A100 GPU.

\subsection{Integration with DeepSpeed ZeRO}

ASO integrates with DeepSpeed ZeRO-3 by hooking the \texttt{backward()} phase. The
semantic score $s_i$ is computed during the forward pass and stored alongside activations.
During backward, the gradient accumulation buffer is modified in-place according to
Equation~\ref{eq:projected_gradient} before the all-reduce collective. This requires
no changes to the optimizer or communication schedule.

\subsection{Computational Overhead}

\begin{table}[H]
\centering
\caption{Computational overhead of ASO relative to baseline training.}
\begin{tabular}{lrrr}
\toprule
\textbf{Component} & \textbf{FLOPs added} & \textbf{Memory (MB)} & \textbf{Wall-clock \%} \\
\midrule
Anchor similarity & $O(Md)$ per token & $M \times d \times 4$ & +1.1\% \\
Projection step   & $O(Rd^2)$ per token & $R \times d \times 4$ & +1.4\% \\
Anchor update     & $O(Md)$ per step & negligible & +0.3\% \\
\midrule
\textbf{Total} & --- & 1024 ($M$=512, $d$=4096) & \textbf{+2.8\%} \\
\bottomrule
\end{tabular}
\label{tab:overhead}
\end{table}

%% -----------------------------------------------------------------------
\section{Experiments}
\label{sec:experiments}
%% -----------------------------------------------------------------------

\subsection{Experimental Setup}

All experiments use the Zen MoDE architecture at three scales: 7B, 32B, and 72B
parameters. Training is performed on clusters of 64--512 H100 SXM5 GPUs with NVLink
interconnects. Baseline training uses identical hyperparameters, differing only in the
absence of the ASO gradient modification.

\subsection{Convergence Speed}

\begin{table}[H]
\centering
\caption{Steps to reach target validation loss on a held-out factual recall benchmark.
Lower is better. ASO consistently reaches the target earlier.}
\begin{tabular}{lrrr}
\toprule
\textbf{Model scale} & \textbf{Baseline (steps)} & \textbf{ASO (steps)} & \textbf{Speedup} \\
\midrule
7B   & 42,000 & 35,200 & $1.19\times$ \\
32B  & 68,000 & 54,600 & $1.25\times$ \\
72B  & 91,000 & 70,800 & $1.28\times$ \\
\bottomrule
\end{tabular}
\label{tab:convergence}
\end{table}

\subsection{Downstream Task Performance}

\begin{table}[H]
\centering
\caption{Benchmark performance comparison. Zen MoDE-72B trained with and without ASO.}
\begin{tabular}{lrrl}
\toprule
\textbf{Benchmark} & \textbf{Baseline} & \textbf{ASO} & \textbf{Improvement} \\
\midrule
TriviaQA (EM)    & 83.4 & 86.1 & +2.7 pp \\
NaturalQuestions & 51.2 & 54.8 & +3.6 pp \\
MMLU (5-shot)    & 84.7 & 85.9 & +1.2 pp \\
HellaSwag        & 87.3 & 87.5 & +0.2 pp \\
GSM8K            & 91.2 & 92.4 & +1.2 pp \\
HumanEval        & 79.3 & 80.1 & +0.8 pp \\
\bottomrule
\end{tabular}
\label{tab:downstream}
\end{table}

ASO improves most on knowledge-retrieval tasks (TriviaQA, NaturalQuestions) and shows
modest but consistent gains on reasoning (GSM8K, HumanEval). Pure language generation
(HellaSwag) shows negligible change, confirming that ASO does not degrade fluency.

\subsection{Continual Learning and Forgetting}

\begin{table}[H]
\centering
\caption{Catastrophic forgetting measured as performance drop on original benchmarks
after one epoch of domain-specific fine-tuning.}
\begin{tabular}{lrr}
\toprule
\textbf{Metric} & \textbf{Baseline} & \textbf{ASO} \\
\midrule
MMLU drop after medical fine-tuning   & $-4.1$ pp & $-1.8$ pp \\
TriviaQA drop after code fine-tuning  & $-6.3$ pp & $-2.9$ pp \\
GSM8K drop after legal fine-tuning    & $-3.7$ pp & $-1.4$ pp \\
\bottomrule
\end{tabular}
\label{tab:forgetting}
\end{table}

The anchor maintenance mechanism in ASO acts as an implicit rehearsal signal,
substantially reducing forgetting by biasing gradients toward the semantic subspace of
the original pretraining distribution.

\subsection{Ablation Studies}

\begin{table}[H]
\centering
\caption{Ablation on ASO components. Each row removes one component. Evaluated on
TriviaQA EM after 50K training steps (Zen MoDE-7B).}
\begin{tabular}{lrr}
\toprule
\textbf{Configuration} & \textbf{TriviaQA EM} & \textbf{$\Delta$ vs. full ASO} \\
\midrule
Full ASO                         & 79.6 & --- \\
No gradient projection           & 77.1 & $-2.5$ \\
No semantic weighting ($\lambda=0$) & 76.8 & $-2.8$ \\
Static anchors (no update)       & 78.2 & $-1.4$ \\
Baseline (no ASO)                & 75.3 & $-4.3$ \\
\bottomrule
\end{tabular}
\label{tab:ablation}
\end{table}

%% -----------------------------------------------------------------------
\section{Discussion}
\label{sec:discussion}
%% -----------------------------------------------------------------------

\subsection{Relation to Importance Sampling}

ASO can be interpreted as an importance sampling correction at the gradient level,
where the importance weight is the embedding-space proximity score $s_i$. Unlike
standard importance sampling, which requires knowing the data distribution, ASO
estimates importance adaptively from the model's own embedding space, making it
self-supervised and free of external annotations.

\subsection{Limitations}

The anchor set introduces a hyper-parameter $M$ (number of anchors) and the anchor
learning rate $\alpha$. In all our experiments, $M = 512$ and $\alpha = 0.01$ worked
well, but optimal values may differ for domain-specific pretraining. Additionally,
the projection step assumes the semantic subspace can be estimated from the top-$R$
principal components of anchor similarities; this approximation degrades for very small
$R$.

\subsection{Future Work}

We are investigating: (1) ASO applied to reinforcement learning from human feedback
(RLHF), where semantic relevance could weight preference pairs; (2) federated ASO
in which anchor sets are maintained across privacy-preserving partitions; (3)
hierarchical anchors that capture both fine- and coarse-grained semantic structure.

%% -----------------------------------------------------------------------
\section{Conclusion}
\label{sec:conclusion}
%% -----------------------------------------------------------------------

Active Semantic Optimization (ASO, HIP-002) is a principled method for biasing
distributed gradient updates toward semantically dense training signal. We have provided
a formal convergence analysis, shown that ASO adds less than 3\% computational overhead,
and demonstrated consistent improvements of 1--4 percentage points on knowledge-intensive
benchmarks at 7B, 32B, and 72B scales, with a 19--28\% reduction in steps to target
validation loss. ASO is available as an open protocol specification at
\url{https://hanzo.ai/hip/002} and a reference implementation in the Zen LM training
codebase.

\section*{Acknowledgements}
The Zen LM Research Team thanks the infrastructure team for cluster support and the
evaluation team for benchmark maintenance.

\begin{thebibliography}{9}
\bibitem{zenlm2025mode}
Zen LM Research Team.
\textit{Zen MoDE: Mixture of Distilled Experts for Scalable Language Models}.
Technical Report v2025.03, Zen LM, 2025.

\bibitem{bengio2009curriculum}
Y. Bengio, J. Louradour, R. Collobert, J. Weston.
\textit{Curriculum Learning}.
ICML, 2009.

\bibitem{kumar2010selfpaced}
M.P. Kumar, B. Packer, D. Koller.
\textit{Self-Paced Learning for Latent Variable Models}.
NeurIPS, 2010.

\bibitem{hacohen2019power}
G. Hacohen, D. Weinshall.
\textit{On the Power of Curriculum Learning in Training Deep Networks}.
ICML, 2019.
\end{thebibliography}

\appendix
\section{Proof of Theorem~\ref{thm:convergence}}

The full convergence proof proceeds in three steps. First, we apply the $L$-smoothness
condition to bound the one-step descent:
\begin{equation}
    \mathcal{L}_{\text{ASO}}(\theta_{t+1}) \leq
    \mathcal{L}_{\text{ASO}}(\theta_t)
    - \eta \langle \nabla \mathcal{L}_{\text{ASO}}(\theta_t), \tilde{G}_t \rangle
    + \frac{L\eta^2}{2}\|\tilde{G}_t\|^2
\end{equation}

Second, we decompose $\tilde{G}_t$ into its signal and noise components, using the
bounded variance assumption A2 and the projection error bound from A4. Third, we
telescope over $T$ steps, optimizing $\eta = \sqrt{1/(LT)}$ to obtain the stated rate.
The additional $\lambda \rho \sigma^2 / L$ term arises from the interaction of
projection error with the semantic weighting factor. $\square$

\end{document}