MOSAIC: Adaptive Inter-layer Composition for Efficient Heterogeneous Vision-Language Models
MOSAIC employs multi-objective MIP to automatically search heterogeneous Transformer architectures, combining linear, sparse, and low-rank mechanisms, boosting efficiency by over 2.5×.
Key Findings
Methodology
MOSAIC defines a unified search space incorporating diverse mechanisms—full, linear, sparse, low-rank attention, and FFN scaling. It employs blockwise local distillation to initialize candidate modules aligned with teacher outputs, then evaluates their capabilities via multiple metrics (perplexity, KL divergence, benchmark scores). The core optimization formulates as a multi-objective mixed-integer programming (MIP) problem, balancing model performance and hardware latency constraints. Solving with Tchebycheff scalarization yields Pareto-optimal architectures. A two-stage parameter recovery—global off-policy distillation stabilizes internal representations, followed by dual-teacher on-policy distillation—restores performance post-structural transition, leveraging a 235B oracle and the original 4B teacher.
Key Results
- The resulting MOSAIC-4B matches the baseline Qwen3-VL-4B-Instruct performance across multiple benchmarks, with less than 2% training cost of the original. It achieves 1.76× speedup in prefill and 2.54× in decoding, significantly improving inference efficiency.
- The architecture search outperforms static manual designs, demonstrating the effectiveness of multi-objective optimization in hardware-constrained environments.
- Ablation studies confirm that integrating multiple mechanisms and capability scoring is crucial for optimal performance and efficiency balance.
Significance
This work advances automated neural architecture search by integrating multi-mechanism heterogeneity with hardware-aware multi-objective optimization. It addresses the longstanding challenge of manually designing static heterogeneous models, providing a scalable framework for deploying high-performance multimodal models on diverse hardware platforms. The approach reduces training costs and inference latency, facilitating real-world applications such as autonomous systems, intelligent assistants, and edge AI. Its ability to adapt architectures automatically to hardware constraints marks a significant step toward democratizing efficient multimodal AI deployment.
Technical Contribution
The paper introduces a novel multi-objective MIP formulation for heterogeneous architecture search, incorporating mechanisms like linear attention, sparse attention, low-rank operators, and FFN scaling. It innovates with blockwise local distillation for initialization and a two-stage parameter recovery strategy that combines off-policy and on-policy distillation with large-capacity teachers. These contributions enable efficient exploration of a vast search space, balancing multiple performance and hardware metrics, and achieving high-quality models with minimal training costs.
Novelty
This is the first work to formulate heterogeneous Transformer architecture search as a multi-objective MIP problem tailored for vision-language models, integrating diverse mechanisms into a unified search space. Unlike prior NAS methods focusing on size or single metrics, this approach simultaneously optimizes multiple objectives—performance, latency, and stability—under hardware constraints, representing a significant innovation in automated model design.
Limitations
- The mechanism pool is predefined, potentially limiting exploration of novel operators or configurations. The search process, while efficient, still requires significant computation, especially for larger models or more mechanisms. Validation across diverse hardware platforms remains limited, necessitating further generalization.
Future Work
Future directions include expanding mechanism pools with emerging operators, integrating end-to-end training with architecture search, and exploring adaptive mechanisms for different hardware environments. Enhancing search efficiency and robustness, as well as deploying models on edge devices with real-time constraints, will further broaden the practical impact of this framework.
AI Executive Summary
In recent years, vision-language models (VLMs) have become fundamental in multimodal AI applications, from autonomous vehicles to intelligent assistants. Traditional homogeneous Transformer architectures, while powerful, face scalability issues due to the quadratic complexity of self-attention, especially in processing long sequences or deploying on resource-constrained hardware. To overcome these limitations, researchers have explored heterogeneous architectures that interleave efficient mechanisms like linear, sparse, and low-rank attention modules, achieving better performance and inference speed. However, manually designing and tuning these hybrid structures is laborious and suboptimal, often tailored to specific hardware without general adaptability.
This paper introduces MOSAIC, a novel automated framework that employs multi-objective mixed-integer programming (MIP) to search for optimal heterogeneous Transformer architectures tailored for hardware constraints. The core idea is to define a comprehensive search space that includes various attention mechanisms and FFN configurations, enabling per-layer customization. The process begins with blockwise local distillation, aligning candidate modules with teacher outputs, followed by multi-capability scoring to evaluate their performance across multiple metrics. The multi-objective optimization balances model accuracy, latency, and stability, producing Pareto-optimal architectures that meet strict hardware latency budgets.
To address performance degradation from structural transitions, the authors propose a two-stage parameter recovery strategy. First, global off-policy distillation stabilizes internal representations; then, dual-teacher on-policy distillation leverages a large oracle model (235B parameters) for knowledge expansion and the original 4B teacher for stability. This approach ensures the final model retains high performance while benefiting from heterogeneous structures.
Experimental results demonstrate that the derived MOSAIC-4B model matches the performance of the baseline Qwen3-VL-4B-Instruct across multiple benchmarks, with training costs reduced to less than 2%. Moreover, inference efficiency improves dramatically, with prefill speedup of 1.76× and decoding acceleration of 2.54×. These findings highlight the potential of automated, hardware-aware architecture search in deploying efficient, high-capacity multimodal models.
Despite its success, the framework relies on predefined mechanisms and still requires significant computational resources for large models. Future work will focus on expanding mechanism diversity, integrating end-to-end training, and deploying on diverse hardware platforms, pushing the boundaries of scalable multimodal AI.
Deep Dive
Plain Language Accessible to non-experts
想象你在厨房里做菜，每次都用一样的锅和调料，虽然简单，但效率不高，也不能做出多样的菜。现在，假设你有一个聪明的厨房助手，它能根据不同的菜肴自动选择最合适的锅、调料和烹饪时间。这样，你就能快速做出各种美味菜肴，而且每次都能用最合适的工具。这就像这篇论文的内容——它开发了一种智能系统，可以自动帮模型选择不同的“工具”和“调料”，让模型在不同硬件环境下都能快速又准确地工作。这个系统会像厨师一样，考虑每个任务的需要，自动调整模型的结构和参数，确保既省时又效果好。这种自动化的“厨房助手”让复杂的模型变得更聪明、更灵活，也更容易在各种设备上使用。
ELI14 Explained like you're 14
想象你在玩一个超级复杂的拼图游戏，有很多不同的拼图块可以用，但每个块的形状和颜色都不一样。你想用最少的时间拼出最漂亮的图案，但每次用不同的拼图块可能会让拼图变得不稳定或者不漂亮。这个论文就像发明了一个聪明的拼图助手，它能自动帮你选择最合适的拼图块组合，让你既能快速拼好，又能拼出最漂亮的图案。它会考虑每个拼图块的特点，还会根据你的目标（比如速度和美观）来调整选择。这样，你不用自己试来试去，就能得到最棒的拼图方案。这就像给模型装上了一个聪明的“拼图助手”，让它在不同的硬件上都能表现得又快又准！
Abstract
Vision-Language Models (VLMs) have achieved success using homogeneous Transformers to process multimedia data. Recent studies show that heterogeneous structures interleaving efficient mechanisms, like linear attention, improve both performance and inference latency over homogeneous designs. However, these efforts rely on handcrafted static mixing patterns, which are sub-optimal and difficult to adapt to specific hardware. To bridge this gap, we propose Multi-Objective Search for Adaptive Inter-layer Composition (MOSAIC), a hardware-aware search method that automatically transforms homogeneous models into optimized heterogeneous architectures. MOSAIC integrates diverse efficiency mechanisms--including linear, sparse, and low-rank operators--into a unified search space. By formulating the selection as a multi-objective Mixed Integer Programming (MIP) problem, our method identifies optimal configurations that maximize downstream performance under strict hardware latency constraints. To mitigate performance degradation from structural transitions, we introduce a two-stage parameter recovery process: global off-policy distillation to stabilize internal representations, followed by a dual-teacher on-policy distillation leveraging a 235B oracle for knowledge expansion and the original 4B teacher for distributional stability. We validate MOSAIC through MOSAIC-4B, derived from Qwen3-VL-4B-Instruct. Results demonstrate that MOSAIC-4B matches the baseline's performance across multiple benchmarks while requiring less than 2% of the original training cost. Furthermore, it substantially improves inference efficiency, achieving 1.76x prefilling and 2.54x decoding speedups.