💻 computer science

Understanding Pruning Regimes in Vision-Language Models Through Domain-Aware Layer Selection

This paper introduces a domain-aware layer selection method for pruning Vision-Language Models that leverages activation similarity to identify critical decoder layers, revealing a three-regime pruning structure where domain-specific rankings ensure stability at low budgets while matching structural baselines at higher budgets.

Original authors: Saeed Khaki, Nima Safaei, Kamal Ginotra

Published 2026-03-24

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Saeed Khaki, Nima Safaei, Kamal Ginotra

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a brilliant, multi-talented assistant named Vision-Language Model (VLM). This assistant can look at a photo of a math problem, a chart, or a messy room, and then answer questions about it. But there's a catch: this assistant is huge. It has a massive brain (a deep neural network) with hundreds of layers of thinking steps. While powerful, this size makes it slow and expensive to run on regular devices.

The paper asks a simple question: "Can we fire some of the assistant's thinking steps to make it faster, without making it forget how to do math or understand pictures?"

Here is the breakdown of their findings using everyday analogies.

1. The Problem: The "Over-Engineered" Brain

Think of the model's brain as a long assembly line in a factory. Each station (or "layer") takes the product (the image and text) and does a little bit of work before passing it to the next station.

The Issue: The factory is so big that many stations are doing very little actual work. Some stations just pass the package along without changing it much.
The Goal: We want to shut down the lazy stations to speed up the factory. But if we shut down the wrong station, the whole product might break.

2. The Old Way vs. The New Way

Previously, people tried to fire stations in two main ways:

Random Firing: "Let's just close every 10th station." (This is risky; you might close the one that actually matters).
Similarity Firing: "Let's close stations that look exactly like the one before them." (This is better, but it doesn't care what the factory is making).

The New Approach (Domain-Aware Pruning):
The authors realized that the factory makes different products for different customers.

Customer A (Math): Needs deep, complex reasoning.
Customer B (General): Needs quick, broad understanding (like "what color is the car?").

They proposed a new rule: "Don't just look at the station; look at what it's doing for this specific customer."

They tested the assistant with two types of questions:

Math questions: "Solve this algebra problem in the image."
General questions: "What is in this picture?"

They measured how much each station changed the information.

If a station barely changed the math answer, it was redundant for math.
If a station barely changed the general description, it was redundant for general tasks.

3. The Three "Regimes" (The Three Stages of Cutting)

The paper discovered that how you fire workers depends entirely on how many you fire. It's like trimming a hedge:

Stage 1: The "Snip" (Low Budget - 10% cut)
- Analogy: You are trimming a few stray branches.
- What happens: It matters exactly which branch you cut. If you cut the one holding the fruit, the tree suffers.
- Result: The authors' "Math-Aware" method wins here. By knowing which branches are useless for math, they can trim the tree perfectly without hurting the fruit. Random cutting hurts performance immediately.
Stage 2: The "Thinning" (Medium Budget - 25% cut)
- Analogy: You are thinning out the middle of the hedge.
- What happens: The tree is getting a bit sparse. It doesn't matter as much which specific branch you cut anymore because the damage is starting to pile up. All methods start to look similar because the tree is just getting weaker.
Stage 3: The "Lopping" (High Budget - 40% cut)
- Analogy: You are chopping off huge chunks of the hedge.
- What happens: The tree is in danger of collapsing. Now, the most important thing isn't which branch you cut, but keeping the structure intact. You need to leave gaps so the remaining branches can still reach the sun.
- Result: A method called "Interlace" (which cuts branches evenly spaced out) wins here. It keeps the tree standing even when it's very short. The "Math-Aware" method loses its advantage because the tree is just too short to function well, regardless of which branches were cut.

4. The "Pareto Frontier" (The Perfect Balance)

The authors found a "sweet spot" strategy.

If you only care about Math, cut the layers that are useless for Math.
If you only care about General Vision, cut the layers useless for General Vision.
The Magic Mix: They found that if you cut based on a mix (mostly General, a little Math), you get the best of both worlds. It's like firing a worker who is bad at math but good at general tasks, while keeping the ones who are good at math. This keeps the assistant smart at math but still good at describing pictures.

5. The "Recovery" Step

After firing the workers, the factory is a bit shaky. The authors added a "Re-Training" phase (Supervised Fine-Tuning).

Analogy: After you fire half the staff, you give the remaining staff a quick 1-day training camp to learn how to work together in the new, smaller factory.
This step was crucial. Without it, the model would be confused. With it, the model recovered most of its lost smarts.

Summary: What Does This Mean for You?

This paper gives us a recipe for making AI assistants faster and cheaper without making them "dumb."

Don't cut randomly. Know what the AI is being asked to do.
Be careful at first. When making small cuts, be very precise about which parts you remove.
Keep the structure. When making big cuts, make sure you don't chop the whole middle out; leave a spaced-out skeleton so the AI can still "think."
Train it again. Always give the AI a quick refresher course after you shrink it.

By using these "Domain-Aware" rules, we can shrink these giant AI models so they can run on your phone or laptop, while still being smart enough to help you solve math problems or read charts.

1. Problem Statement

Vision-Language Models (VLMs) have achieved state-of-the-art performance in multimodal tasks but suffer from high computational costs and latency due to their deep transformer architectures. While pruning (removing redundant components) is a standard compression technique, existing methods often rely on aggregate metrics or text-only heuristics that fail to account for the unique demands of multimodal reasoning.

The Core Challenge:

Domain Sensitivity: VLMs must handle diverse tasks ranging from simple captioning to complex multi-step mathematical reasoning. The paper argues that "redundancy" is not uniform; a layer critical for visual grounding might be redundant for math reasoning, and vice versa.
The "Regime" Gap: It is unclear how depth pruning behaves across different pruning budgets (low vs. high) and whether the specific choice of removed layers matters more at certain stages of compression.
Mathematical Reasoning: Math problems in VLMs require a tight coupling of perception (diagrams, symbols) and language reasoning. Small structural changes can disproportionately degrade math performance even if general vision-language scores remain stable.

2. Methodology

The authors propose a Domain-Aware Decoder Layer Pruning framework that selects layers for removal based on how much they transform input representations for specific domains (Math vs. Non-Math).

A. Activation Similarity Metric

Instead of using static weight magnitudes, the method measures input-output activation similarity (cosine similarity) for each decoder layer $\ell$ .

Process: For a given domain (Math or Non-Math), the model processes a diverse set of probes (e.g., Chain-of-Thought math prompts vs. captioning prompts).
Redundancy Score: If a layer's output activation is highly similar to its input activation, the layer is deemed functionally redundant for that domain.
Aggregation: Scores are aggregated across multiple subtasks and normalized (z-score) to ensure comparability across the model depth.

B. Ranking Strategies

The authors introduce three ranking variants to determine which layers to prune:

Math-Aware: Prunes layers with high similarity scores specifically on math prompts.
Non-Math-Aware: Prunes layers with high similarity scores on general vision-language prompts.
Mixed (Hybrid): A weighted combination ( $R_\ell = \alpha \hat{S}_{nonmath} + (1-\alpha) \hat{S}_{math}$ ). The authors empirically found $\alpha=0.7$ optimal, prioritizing general capabilities (OCR, grounding) while preserving math reasoning.

C. Pruning Protocol

Structured Removal: Entire decoder blocks are removed.
Constraints: The first and last decoder layers are protected to maintain multimodal fusion and final token prediction stability.
Post-Pruning Stabilization: A brief Supervised Fine-Tuning (SFT) stage is applied to all pruned models using a mix of math and general instruction data to recover performance.

3. Key Contributions

Domain-Aware Pruning Method: A novel approach using input-output activation similarity to identify domain-specific redundancy, yielding math-aware, non-math-aware, and mixed ranking criteria.
Empirical Regime Analysis: The paper identifies a consistent three-regime structure in VLM depth pruning:
- Ranking-Sensitive Regime (Low Budget, <15%): Performance is highly sensitive to which layers are removed. Domain-aware selection is critical here.
- Transition Regime (Moderate Budget, 15–32%): Methods begin to converge as structural damage accumulates.
- Structure-Dominated Regime (High Budget, >32%): Structural continuity becomes the limiting factor. Strategies that enforce spacing between deletions (avoiding contiguous holes) outperform pure redundancy ranking.
Comprehensive Benchmarking: Evaluation on two state-of-the-art models (Qwen3-VL-2B and 4B) across math-specific benchmarks (Snapask, NuminaMath) and general multimodal benchmarks (ChartQA, ScienceQA, RealWorldQA, VStar).

4. Experimental Results

A. Performance Across Regimes

Low Budget (10%): Domain-aware methods (especially the Mixed variant) significantly outperform baselines. They preserve the "Pareto frontier," balancing math retention and general vision-language capability.
- Example: On Qwen3-VL-2B at 10% pruning, the Mixed method achieved a 24.88% accuracy on Snapask, whereas Random pruning dropped to 17.09%.
Moderate Budget (25%): All methods suffer significant degradation, but domain-aware methods remain slightly more stable.
High Budget (40%): Structural continuity dominates. The Interlace baseline (which enforces spaced deletions) often outperforms pure redundancy ranking because it prevents the collapse of residual flow caused by removing adjacent layers.

B. Trade-offs (Math vs. General VL)

Math-Aware Pruning: Shifts performance toward better math retention but risks degrading general visual grounding.
Non-Math-Aware Pruning: Preserves general capabilities (OCR, spatial reasoning) but degrades math reasoning.
Mixed Strategy: Provides the best overall balance, lying near the Pareto frontier. It prevents the "brittleness" seen in random or purely structure-agnostic pruning.

C. Baseline Comparisons

CKA (Centered Kernel Alignment): A representation-similarity baseline that tends to prune early layers aggressively, leading to poor performance on tasks requiring deep reasoning.
Random Pruning: Consistently the worst performer, highlighting that redundancy is not uniformly distributed.
Interlace: Performs well at high budgets due to its spacing constraint but is less effective than domain-aware methods at low budgets where specific layer selection matters most.

5. Significance and Implications

Interpretability: The work provides a clear, interpretable mechanism for understanding why certain layers are critical for specific tasks (e.g., math reasoning relies heavily on mid-to-late layers, while grounding relies on early layers).
Practical Deployment: Offers a practical recipe for compressing VLMs without sacrificing critical domain capabilities. By using a "Mixed" ranking, practitioners can reduce model depth by 10–25% with minimal impact on math or general reasoning.
Theoretical Insight: The identification of the three-regime structure is a major theoretical contribution. It suggests that pruning strategies should be adaptive: use domain-aware selection for light compression and structure-aware spacing for heavy compression.
Future Directions: The findings suggest that future pruning research for multimodal models must move beyond aggregate metrics and consider domain-specific activation patterns and structural continuity simultaneously.

In summary, this paper demonstrates that not all layers are created equal in VLMs. By aligning pruning decisions with domain-specific activation patterns, one can achieve efficient compression while preserving the complex reasoning capabilities required for mathematical and scientific tasks.