MOSAIC: Adaptive Inter-layer Composition for Efficient Heterogeneous Vision-Language Models¶

🕒 Published (v1): 2026-07-10 01:31 UTC · Source: Arxiv · link

Why this paper was selected

Heterogeneous VLM with linear attention interleaving; efficiency + performance; actionable architecture direction

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

TL;DR¶

MOSAIC is a hardware-aware neural architecture search method that automatically transforms homogeneous VLMs into heterogeneous architectures by mixing full, linear, sparse, and low-rank attention operators per layer. It formulates layer selection as a multi-objective Mixed Integer Programming problem under strict latency constraints, then recovers performance via two-stage distillation. Applied to Qwen3-VL-4B-Instruct, the resulting MOSAIC-4B achieves 2.54× decoding acceleration and 1.76× prefilling speedup at matched accuracy using less than 2% of the original training cost.

Problem¶

Heterogeneous VLMs that interleave efficient attention mechanisms (linear, sparse, SSM) with standard attention currently rely on manually crafted, static mixing patterns (e.g., periodic interleaving every $k$ layers). These heuristics are sub-optimal, ignore depth-dependent representational demands, and cannot adapt to varying hardware latency profiles without expensive end-to-end retraining.

Method¶

Search Space. Each of the $L=36$ transformer layers independently selects an attention subblock from: GQA (original), KDA (linear), GDN (linear), MLA (low-rank KV), SWA (sparse); and an FFN subblock scaled to $\{0\%, \ldots, 100\%\}$ of the original intermediate dimension. Total space: $35^{36}$.

Blockwise Local Distillation (BLD). Each candidate block is pre-trained to minimize $\mathcal{L}_\text{align} = \|\mathbf{o}_t - \mathbf{o}_c\|^2 / \|\mathbf{o}_t\|^2$ against the corresponding teacher layer, using only 80k samples. This is parallelizable and hardware-agnostic.

Multi-Capability Scoring. Each candidate block receives per-ability scores: $$S^k_{l,i} = (-1)^{\sigma_k} \cdot \frac{P^k_{l,i} - M^k}{M^k}$$ across four objectives: KL divergence, perplexity, LLM benchmark, VLM benchmark.

Multi-Objective MIP. Architecture selection is cast as: $$\max_{X \in \mathcal{B}} \left\{ F^1(X), \ldots, F^m(X) \right\} \quad \text{s.t.} \sum_{i,l} \text{runtime}(l,i)\, x_{l,i} \leq \text{Target}$$ solved via Tchebycheff scalarization, converted to a linear MIP with auxiliary variable $C$.

Two-Stage Parameter Recovery. 1. Global off-policy distillation: student is trained against Qwen3-VL-4B-Instruct via $\mathcal{L}_\text{recover} = \mathcal{L}_\text{KL}(\theta; D) + \alpha \mathcal{L}_\text{CE}(\theta; D)$ using symmetric KL divergence on text + image-video-text data. 2. Dual-teacher on-policy distillation: a 235B oracle teacher (Qwen3-VL-235B) provides SFT targets for knowledge expansion; the original 4B teacher provides on-policy KL supervision on student-generated tokens to prevent drift. Objective: $\mathcal{L}_\text{align} = \mathcal{L}_\text{KL}(\theta; D') + \beta \mathcal{L}_\text{CE}(\theta; D_o)$.

Key Contributions¶

First automated search for optimal heterogeneous VLM configurations (per-layer mechanism type as primary search variable).
Unified search space spanning linear (KDA, GDN), sparse (SWA), low-rank (MLA), and full (GQA) attention, plus FFN width scaling.
Multi-objective MIP formulation with hardware-measured runtimes enabling hardware-specific Pareto-optimal solutions.
Progressive two-stage distillation separating representation stabilization (off-policy, 4B teacher) from capability expansion (on-policy, 235B oracle).
MOSAIC-4B: a publicly released heterogeneous derivative of Qwen3-VL-4B-Instruct.

Results¶

Decoding acceleration (TPOT): 2.54× over Qwen3-VL-4B-Instruct baseline.
Prefilling speedup (TTFT): 1.76× at 96k context; 2.72× at shorter prefill lengths observed in ablation.
Accuracy: Matches Qwen3-VL-4B-Instruct across multiple image understanding benchmarks.
Training cost: Less than 2% of original model training cost.
MOSAIC-4B outperforms comparably-sized models (InternVL3.5-4B, Molmo2-4B, Qwen3-VL-2B-Instruct, InternVL3.5-2B, SmolVLM2-2.2B) on the average performance vs. decoding acceleration Pareto frontier (Fig. 1).

Limitations¶

Search space is constrained to architectures with the same hidden dimension as the parent model; depth reduction or cross-layer restructuring is not explored.
BLD initializes blocks in isolation, creating inter-block compatibility gaps that require an additional recovery stage.
The MIP runtime profiles are measured at fixed batch size (1) and specific context lengths (96k prefill, 256 decode); optimality may not transfer to other deployment configurations.
Validated on a single parent model (Qwen3-VL-4B-Instruct); generalizability to other model families or scales is undemonstrated.
Performance of purely linear attention structures remains constrained by finite-state capacity on very long sequences, which may limit gains in certain long-context tasks.

Relevance to Vision-Language Models¶

MOSAIC directly addresses a core deployment bottleneck for VLMs: the quadratic complexity of standard attention under long-context or edge-device constraints, which is especially acute in multimodal settings where visual tokens inflate sequence length. By automating heterogeneous architecture search with hardware-in-the-loop latency constraints, it advances a practical path for adapting existing high-quality VLMs (rather than training from scratch) to specific deployment targets. The dual-teacher distillation strategy, using a much larger oracle to expand capability while the original teacher anchors stability, is a transferable training recipe for model compression in the VLM space. This work extends the line of hybrid-architecture research (Jamba, Kimi-Linear, InfiniteVL) from hand-designed to data-driven composition, which is likely to become the standard approach as the heterogeneous VLM design space grows.