papers/zip-002-zen-reranker/sections/04-zen-reranker-architecture.tex

\section{Zen-Reranker Architecture}


\subsection{Model Design}

Zen-Reranker extends Qwen3-Embedding-8B with a specialized projection head:

\begin{equation}
\mathbf{h} \in \mathbb{R}^{8192} \xrightarrow{\text{Projection Head}} \mathbf{e} \in \mathbb{R}^{7680}
\end{equation}

\textbf{Projection Head Architecture:}
\begin{align}
\mathbf{z}_1 &= \text{GELU}(\mathbf{W}_1 \mathbf{h} + \mathbf{b}_1), \quad \mathbf{W}_1 \in \mathbb{R}^{6144 \times 8192} \\
\mathbf{z}_2 &= \text{LayerNorm}(\mathbf{z}_1) \\
\mathbf{z}_3 &= \mathbf{W}_2 \mathbf{z}_2 + \mathbf{b}_2, \quad \mathbf{W}_2 \in \mathbb{R}^{7680 \times 6144} \\
\mathbf{e} &= \frac{\text{LayerNorm}(\mathbf{z}_3)}{||\text{LayerNorm}(\mathbf{z}_3)||_2} \quad \text{(L2 normalize)}
\end{align}

\textbf{Design Rationale:}
\begin{itemize}
\item \textbf{Two-layer MLP:} Sufficient capacity to learn 8192$\to$7680 mapping with non-linearity
\item \textbf{Intermediate dimension 6144:} Balances expressiveness vs. parameter efficiency ($\frac{3}{4} \times 8192$)
\item \textbf{GELU activation:} Smooth gradients, better than ReLU for embedding learning
\item \textbf{LayerNorm after each layer:} Stabilizes training, reduces internal covariate shift
\item \textbf{L2 normalization:} Ensures unit hypersphere embeddings, critical for cosine similarity
\end{itemize}

\textbf{Parameter Count:}
\begin{align}
\text{Params} &= (8192 \times 6144) + 6144 + (6144 \times 7680) + 7680 \\
&= 50,331,648 + 6,144 + 47,185,920 + 7,680 \\
&= \mathbf{97,531,392} \text{ parameters (97.5M)}
\end{align}

With base Qwen3-8B (8.2B params), total model size: \textbf{8.3B parameters}.

\subsection{Training Protocol}

We adopt a three-stage training strategy, each optimizing different objectives:

\subsubsection{Stage 1: Projection Expansion (18 hours, 8$\times$ H100)}

\textbf{Objective:} Learn projection head to match Qwen3-8B's semantic space in 7680 dimensions.

\textbf{Dataset:} 100M text pairs from MS MARCO~\cite{nguyen2016msmarco} and Natural Language Inference~\cite{bowman2015snli}.

\textbf{Loss Function:}
\begin{equation}
\mathcal{L}_{\text{expansion}} = \frac{1}{N} \sum_{i=1}^{N} ||\mathbf{e}_i^{7680} - \text{Pad}(\mathbf{e}_i^{4096})||_2^2
\end{equation}

where $\text{Pad}(\mathbf{e}^{4096})$ zero-pads Qwen3's 4096-dim embedding to 7680-dim.

\textbf{Hyperparameters:}
\begin{itemize}
\item Learning rate: $5 \times 10^{-4}$ with linear warmup (1,000 steps)
\item Optimizer: AdamW ($\beta_1=0.9, \beta_2=0.999, \epsilon=10^{-8}$)
\item Batch size: 2,048 (256 per GPU $\times$ 8 GPUs)
\item Weight decay: 0.01
\item Gradient clipping: 1.0
\end{itemize}

\textbf{Result:} Projection head approximates Qwen3's semantic space with 96.3\% cosine similarity on held-out validation set.

\subsubsection{Stage 2: Reranking Fine-tuning (12 hours, 8$\times$ H100)}

\textbf{Objective:} Optimize for retrieval tasks, improving ranking quality.

\textbf{Dataset:} TREC-COVID~\cite{voorhees2020trec}, MS MARCO Passage, BEIR~\cite{thakur2021beir}.

\textbf{Loss Function:} Contrastive loss with hard negatives~\cite{karpukhin2020dpr}:
\begin{equation}
\mathcal{L}_{\text{contrast}} = -\log \frac{\exp(\mathbf{q} \cdot \mathbf{p}^+ / \tau)}{\exp(\mathbf{q} \cdot \mathbf{p}^+ / \tau) + \sum_{j=1}^{K} \exp(\mathbf{q} \cdot \mathbf{p}_j^- / \tau)}
\end{equation}

where $\mathbf{q}$ is query embedding, $\mathbf{p}^+$ is positive passage, $\mathbf{p}_j^-$ are hard negatives, and $\tau=0.05$ is temperature.

\textbf{Hyperparameters:}
\begin{itemize}
\item Learning rate: $1 \times 10^{-5}$ (10$\times$ lower than Stage 1)
\item Batch size: 1,024 (128 per GPU $\times$ 8 GPUs)
\item Hard negatives per query: $K=7$
\item Training steps: 50,000
\end{itemize}

\textbf{Result:} MTEB Retrieval score improves from 61.3 to 62.7 (+1.4 points).

\subsubsection{Stage 3: DSO Optimization (24 hours, 8$\times$ H100)}

\textbf{Objective:} Co-optimize for BitDelta compression, Byzantine robustness, and semantic diversity.

\textbf{Dataset:} 5M synthetic DSO scenarios generated via prompt engineering:
\begin{itemize}
\item 2M math problem-solving experiences
\item 2M coding task experiences  
\item 1M general reasoning experiences
\end{itemize}

\textbf{Loss Function:} Multi-objective optimization:
\begin{equation}
\mathcal{L}_{\text{DSO}} = \lambda_1 \mathcal{L}_{\text{compress}} + \lambda_2 \mathcal{L}_{\text{robust}} + \lambda_3 \mathcal{L}_{\text{diversity}}
\end{equation}

\textbf{Compression Loss:} Penalize high entropy in delta encoding:
\begin{equation}
\mathcal{L}_{\text{compress}} = \frac{1}{N} \sum_{i=1}^{N} H(\Delta \mathbf{e}_i), \quad \Delta \mathbf{e}_i = [\mathbf{e}_i[j] - \mathbf{e}_i[j-1]]_{j=1}^{7680}
\end{equation}

where $H(\cdot)$ is empirical entropy of deltas.

\textbf{Robustness Loss:} Minimize distance between clean median and median under simulated attack:
\begin{equation}
\mathcal{L}_{\text{robust}} = \frac{1}{N} \sum_{i=1}^{N} ||\text{Median}(\{\mathbf{e}_i^{(j)}\}_{j=1}^{n}) - \text{Median}(\{\mathbf{e}_i^{(j)}\}_{j=1}^{n} \cup \{\mathbf{a}_i^{(k)}\}_{k=1}^{m})||_2^2
\end{equation}

where $\{\mathbf{a}_i^{(k)}\}$ are adversarial embeddings (random noise or adversarially perturbed).

\textbf{Diversity Loss:} Encourage coverage of semantic space:
\begin{equation}
\mathcal{L}_{\text{diversity}} = -\frac{1}{N(N-1)} \sum_{i \neq j} ||\mathbf{e}_i - \mathbf{e}_j||_2^2
\end{equation}

\textbf{Hyperparameters:}
\begin{itemize}
\item Loss weights: $\lambda_1=0.3, \lambda_2=0.5, \lambda_3=0.2$
\item Learning rate: $5 \times 10^{-6}$
\item Batch size: 512 (64 per GPU $\times$ 8 GPUs)
\item Adversarial ratio: 30\% Byzantine nodes in robustness simulation
\end{itemize}

\textbf{Result:}
\begin{itemize}
\item BitDelta compression: 31.87$\times$ (vs 29.9$\times$ without this stage)
\item Accuracy under 30\% attack: 92.1\% (vs 87.3\% baseline)
\item Semantic diversity (avg pairwise distance): 1.87 (vs 1.62 baseline)
\end{itemize}

\subsection{Training Cost Analysis}

\begin{table}[h]
\centering
\caption{Training cost breakdown (8$\times$ H100 80GB @ \$2/GPU-hour)}
\begin{tabular}{lrrr}
\toprule
\textbf{Stage} & \textbf{Duration (hours)} & \textbf{Cost} & \textbf{Percentage} \\
\midrule
Stage 1: Projection & 18 & \$3,600 & 33.3\% \\
Stage 2: Reranking & 12 & \$2,400 & 22.2\% \\
Stage 3: DSO & 24 & \$4,800 & 44.4\% \\
\midrule
\textbf{Total} & \textbf{54} & \textbf{\$10,800} & \textbf{100\%} \\
\bottomrule
\end{tabular}
\end{table}

\textbf{Cost Comparison:}
\begin{itemize}
\item Training from scratch (8B model): \$50,000+
\item Full fine-tuning Qwen3-8B: \$25,000
\item \textbf{Our three-stage approach: \$10,800 (78\% savings vs. scratch)}
\end{itemize}