papers/zen-multimodal-architecture.tex

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

\title{\textbf{Zen Multimodal Architecture: Unified Vision-Language-Audio}\\
\large Technical Report v2025.03}
\author{Antje Worring, Zach Kelling \\ Zen LM Research Team\\
\texttt{[email protected]}}
\date{March 2025}

\begin{document}
\maketitle

\begin{abstract}
We present the Zen multimodal architecture, a unified framework for processing and
generating text, images, and audio within a single model. Our approach combines
a late-fusion cross-attention mechanism with modal-specific encoders and a universal
vocabulary that represents all modalities as discrete tokens. Key innovations include
adaptive modal dropout during training, which teaches the language backbone to operate
effectively with any subset of modalities present, and a learned modal routing mechanism
that dynamically allocates compute based on input modality complexity. The architecture
achieves state-of-the-art results on MMBench (82.4), VideoQA (78.1\%), AudioCaps
retrieval (R@1 = 64.3\%), and cross-modal retrieval (MSCOCO R@1 = 87.2\%) while
maintaining text-only performance within 0.3 points of the text-specialized baseline.
\end{abstract}

\section{Introduction}

Human cognition integrates perception across modalities seamlessly. Language models
operating on text alone miss the rich structure available in images, video, and audio.
Multimodal models must address three challenges:

\begin{enumerate}
  \item \textbf{Alignment}: Mapping each modality into a shared semantic space
  \item \textbf{Fusion}: Combining cross-modal information for joint reasoning
  \item \textbf{Generation}: Producing outputs in any modality
\end{enumerate}

Early multimodal models concatenated visual tokens to text context (ViLBERT, early
GPT-4V style approaches). This is limited by context length and fails to capture
fine-grained cross-modal attention patterns. Zen's architecture uses late fusion:
each modality is processed by a specialized encoder, and cross-modal attention is
applied at the language backbone level.

\section{Architecture Overview}

\subsection{Modal Encoders}

\subsubsection{Vision Encoder}

The vision encoder is a ViT-L/14 variant with 307M parameters, trained from scratch
on 2B image-text pairs. It produces patch embeddings of dimension 1024 at
14$\times$14 pixel patches. For a $224\times224$ image, this yields $16\times16=256$
visual tokens.

Dynamic resolution processing supports arbitrary input sizes:
\begin{equation}
N_\text{patches} = \left\lceil H / 14 \right\rceil \times \left\lceil W / 14 \right\rceil
\end{equation}
with 2D rotary positional embeddings that interpolate smoothly to unseen resolutions.

\subsubsection{Audio Encoder}

The audio encoder converts raw waveforms to mel spectrograms and processes them
with a 1D convolutional front-end followed by a transformer:

\begin{equation}
\text{mel}[t, f] = \log\left(\sum_k w[k] \cdot |X[t, k]|^2 + \epsilon\right)
\end{equation}

with 80 mel bins, 10ms frame shift, and 25ms window. The transformer produces
audio tokens at 50Hz (one token per 20ms). Audio sequences up to 30 seconds
are supported natively (1500 tokens).

\subsubsection{Universal Tokenizer}

Text tokens follow the standard BPE vocabulary (150K tokens). Visual and audio
tokens are produced by vector-quantized encoders with codebook size 8192:

\begin{equation}
q(z) = \arg\min_{e \in \mathcal{E}} \|z - e\|_2
\end{equation}

where $\mathcal{E}$ is the learned codebook. Visual codebooks are trained on
image reconstruction; audio codebooks on waveform reconstruction with a GAN loss.
This gives a unified discrete vocabulary of $150\text{K} + 2\times8192 = 166\text{K}$
tokens across all modalities.

\subsection{Cross-Modal Attention Fusion}

The language backbone (Zen-7B) processes the combined sequence of text, visual,
and audio tokens. To prevent text processing from being dominated by the larger
visual/audio token sequences, we use gated cross-modal attention:

\begin{equation}
\text{CrossAttn}(Q_\text{text}, K_\text{vis}, V_\text{vis}) = \text{softmax}\left(\frac{Q_\text{text} K_\text{vis}^\top}{\sqrt{d}}\right) V_\text{vis}
\end{equation}

A gating scalar $\alpha$ interpolates between self-attention output and cross-attention:

\begin{equation}
h = (1 - \alpha) \cdot \text{SelfAttn}(h) + \alpha \cdot \text{CrossAttn}(h, V_\text{vis})
\end{equation}

with $\alpha$ learned per layer via a sigmoid gate initialized to 0. This allows
the model to gradually learn cross-modal integration rather than immediately
disrupting pretrained text representations.

\subsection{Modal Dropout}

During training, we randomly drop input modalities with probability:
\begin{itemize}
  \item Drop vision: 15\%
  \item Drop audio: 20\%
  \item Drop both: 5\%
  \item Keep all: 60\%
\end{itemize}

Modal dropout teaches the model to extract maximum information from available
modalities and prevents over-reliance on any single modality. It also enables
zero-shot transfer to text-only, vision-only, and audio-only inference modes.

\section{Training}

\subsection{Stage 1: Modal Encoder Pretraining}

Vision encoder: trained on 2B image-text pairs (LAION-5B subset) with contrastive
CLIP-style loss for 100K steps.

Audio encoder: trained on 500K audio-text pairs (AudioCaps, AudioSet with captions)
with contrastive loss for 50K steps.

\subsection{Stage 2: Unified Backbone Training}

The pretrained modal encoders are connected to the Zen-7B language backbone. We
freeze encoder weights and train only the cross-attention gates and projection layers
for 50K steps on multimodal instruction data:

\begin{table}[H]
\centering
\begin{tabular}{lrr}
\toprule
\textbf{Data Type} & \textbf{Samples} & \textbf{Mix Weight} \\
\midrule
Image + text Q\&A & 1.2M & 40\% \\
Image captioning & 800K & 20\% \\
Video + text & 400K & 15\% \\
Audio + text & 300K & 10\% \\
Text only & 2.0M & 15\% \\
\bottomrule
\end{tabular}
\caption{Stage 2 multimodal training data composition.}
\end{table}

\subsection{Stage 3: End-to-End Fine-Tuning}

All components are fine-tuned jointly with a lower learning rate ($\eta=10^{-5}$)
for 20K steps. The text-only loss coefficient is upweighted by 2$\times$ to
prevent catastrophic forgetting of text capabilities.

\section{Experiments}

\subsection{Visual Understanding: MMBench}

\begin{table}[H]
\centering
\begin{tabular}{lccccc}
\toprule
\textbf{Model} & \textbf{Overall} & \textbf{LR} & \textbf{AR} & \textbf{RR} & \textbf{FP} \\
\midrule
Zen-7B-Multimodal & \textbf{82.4} & 80.1 & 84.3 & 79.8 & 85.2 \\
Zen-32B-Multimodal & 85.1 & 83.4 & 86.8 & 82.3 & 88.1 \\
\bottomrule
\end{tabular}
\caption{MMBench results. LR=Logical Reasoning, AR=Attribute Recognition, RR=Relation Reasoning, FP=Fine-grained Perception.}
\end{table}

\subsection{Video Understanding: VideoQA}

\begin{table}[H]
\centering
\begin{tabular}{lcccc}
\toprule
\textbf{Model} & \textbf{MSVD-QA} & \textbf{MSRVTT-QA} & \textbf{ActivityNet-QA} & \textbf{Avg.} \\
\midrule
Zen-7B-Multimodal & 79.3\% & 76.8\% & 78.2\% & 78.1\% \\
Zen-32B-Multimodal & 82.1\% & 79.4\% & 81.0\% & 80.8\% \\
\bottomrule
\end{tabular}
\caption{VideoQA accuracy across benchmarks.}
\end{table}

\subsection{Audio Understanding: AudioCaps}

\begin{table}[H]
\centering
\begin{tabular}{lcccc}
\toprule
\textbf{Model} & \textbf{R@1} & \textbf{R@5} & \textbf{R@10} & \textbf{CIDEr} \\
\midrule
Zen-7B-Multimodal & 64.3 & 88.2 & 94.1 & 82.4 \\
\bottomrule
\end{tabular}
\caption{AudioCaps text-to-audio retrieval and captioning results.}
\end{table}

\subsection{Cross-Modal Retrieval: MSCOCO}

\begin{table}[H]
\centering
\begin{tabular}{lcccccc}
\toprule
\textbf{Model} & \multicolumn{3}{c}{\textbf{Image$\to$Text}} & \multicolumn{3}{c}{\textbf{Text$\to$Image}} \\
& \textbf{R@1} & \textbf{R@5} & \textbf{R@10} & \textbf{R@1} & \textbf{R@5} & \textbf{R@10} \\
\midrule
Zen-7B-Multimodal & 87.2 & 97.8 & 99.1 & 72.4 & 91.3 & 96.2 \\
\bottomrule
\end{tabular}
\caption{MSCOCO 5K test set cross-modal retrieval.}
\end{table}

\subsection{Text-Only Performance Preservation}

\begin{table}[H]
\centering
\begin{tabular}{lcccc}
\toprule
\textbf{Model} & \textbf{MMLU} & \textbf{HumanEval} & \textbf{MATH} & \textbf{MT-Bench} \\
\midrule
Zen-7B (text only) & 85.3 & 78.2 & 67.4 & 8.62 \\
Zen-7B-Multimodal & 85.1 & 78.0 & 67.2 & 8.59 \\
$\Delta$ & $-0.2$ & $-0.2$ & $-0.2$ & $-0.03$ \\
\bottomrule
\end{tabular}
\caption{Text performance of multimodal vs. text-only model. Modal dropout prevents degradation.}
\end{table}

\section{Analysis}

\subsection{Ablation: Cross-Attention Fusion Strategy}

\begin{table}[H]
\centering
\begin{tabular}{lcc}
\toprule
\textbf{Fusion Strategy} & \textbf{MMBench} & \textbf{Text MMLU} \\
\midrule
Prefix concatenation & 76.2 & 83.1 \\
Early fusion (all layers) & 79.8 & 82.4 \\
Late fusion (top 1/3 layers) & 81.3 & 84.7 \\
Gated cross-attention (ours) & \textbf{82.4} & \textbf{85.1} \\
\bottomrule
\end{tabular}
\caption{Fusion strategy ablation. Gated cross-attention best preserves text while gaining vision.}
\end{table}

\subsection{Effect of Modal Dropout Rate}

\begin{table}[H]
\centering
\begin{tabular}{lccc}
\toprule
\textbf{Dropout Rate} & \textbf{MMBench} & \textbf{Text MMLU} & \textbf{Text-Only Inference} \\
\midrule
0\% & 83.1 & 83.4 & 79.2 \\
10\% & 82.8 & 84.6 & 83.1 \\
20\% (default) & 82.4 & 85.1 & 84.9 \\
40\% & 80.9 & 85.2 & 85.3 \\
\bottomrule
\end{tabular}
\caption{Effect of modal dropout on multimodal and text-only performance.}
\end{table}

Higher dropout rates improve text-only inference robustness at a small cost to
peak multimodal performance. The 20\% default balances both objectives.

\section{Conclusion}

The Zen multimodal architecture demonstrates that a unified vision-language-audio
model can achieve strong performance across all modalities with minimal sacrifice
to text-only capability. Gated cross-attention fusion and modal dropout are the
key methodological contributions enabling this balance. The architecture naturally
extends to video (sequential image frames) and future modalities via the universal
discrete tokenizer.

\bibliographystyle{plain}
\begin{thebibliography}{99}
\bibitem{radford2021clip} Radford et al. (2021). Learning Transferable Visual Models From Natural Language Supervision. \textit{ICML}.
\bibitem{dosovitskiy2021vit} Dosovitskiy et al. (2021). An Image is Worth 16x16 Words. \textit{ICLR}.
\bibitem{gemmateam2024gemini} Gemini Team (2024). Gemini: A Family of Highly Capable Multimodal Models. \textit{arXiv:2312.11805}.
\bibitem{liu2023llava} Liu et al. (2023). Visual Instruction Tuning. \textit{NeurIPS}.
\bibitem{li2023blip2} Li et al. (2023). BLIP-2: Bootstrapping Language-Image Pre-training with Frozen Image Encoders and Large Language Models. \textit{ICML}.
\end{thebibliography}

\end{document}