Gemma-Prune: Compressing Gemma 3 4B Vision-Language Model for Mobile Devices

Model Architecture

Gemma 3 4B IT is a vision-language model combining a text decoder with a SigLIP vision encoder and a multimodal projector:

Seven-Stage Compression Pipeline

Each stage targets a distinct source of redundancy. The techniques compose without interference and accumulate to a total savings of 709 MB.

Vocabulary Pruning & BPE Coverage

The first step reduces the embedding matrix from 262K to 144K tokens using a token map indirection scheme:

A critical finding: dictionary-only pruning (80K tokens) causes generation collapse. BPE tokenizers produce merged tokens like “_the”, “_and”, “_is” that don't appear in any dictionary. Missing these tokens fragments generation output completely.

The solution is an ASCII vocabulary scan — collecting all tokens whose surface forms use only ASCII characters. This brings coverage from 80K to 144K tokens, completely resolving quality issues. This establishes a practical lower bound of ~144K tokens for English deployment.

The Prime Factor Problem

SigLIP's intermediate dimension is 4304 = 16 × 269, where 269 is prime. This means the fc2 weight matrix cannot be directly quantized with MLX group sizes (32, 64, 128).

The solution: zero-pad to 4352 = 68 × 64, aligned to group size 64. This is mathematically lossless — appended zero neurons contribute nothing to dot products, preserving layer output exactly.

Vision Encoder Sensitivity

A surprising asymmetry: removing 3 text decoder layers (159 MB) causes only minor quality loss, while removing just 4 vision encoder layers (35 MB) causes complete image understanding failure — pizza misidentified as “skin texture”, hallucinated descriptions throughout.

Resolution as Thermal Management

Reducing resolution from 896×896 to 672×672 pixels exploits quadratic attention scaling:

A 25% linear reduction yields ~68% attention cost savings — more effective for mobile thermal management than layer reduction, which only scales linearly.

MLP Neuron Pruning

Activation profiling across 20 forward passes reveals that layers 14–30 have 60–100% dead neurons (mean activation < 0.5). Layer 29 is nearly 100% dead. Layers 0–13 maintain healthy activation patterns and are protected from pruning.

For each layer, we retain neurons with the highest activation magnitudes, subject to a maximum 25% reduction ratio and quantization group alignment (group_size=64). This removes 188 MB of dead weight.

Results

Text Generation Performance

Image Understanding Performance

Failed Experiments

Incremental Compression Analysis

Swift Inference Engine

The full inference pipeline is implemented in Swift using Apple's MLX framework with the MLXVLM library, requiring no Python dependencies.

Token Map Support

class MLXArrayBox {
    let value: MLXArray
    init(_ value: MLXArray) { self.value = value }
}

func callAsFunction(_ x: MLXArray) -> MLXArray {
    let indices = tokenMap.value.gathered(x.flattened())
    return compactEmbedding(indices)
        .reshaped(shape + [embeddingDimension])
}

Per-Layer Intermediate Sizes

for i in 0..<config.numHiddenLayers {
    let intermediateSize =
        config.perLayerIntermediateSizes?[i]
        ?? config.intermediateSize
    layers.append(TransformerLayer(config, intermediateSize))
}

Quick Start

gemma-cli --model /path/to/compressed-model \
          --prompt "Describe this image" \
          --image photo.jpg \
          --max-tokens 200 \
          --temperature 0.0

Pre-built Models

Three model configurations are available, all on HuggingFace and GitHub:

Citation

@article{atomgradient2026gemmaprune,
  title={Gemma-Prune: A Multi-Stage Compression Pipeline
         for Deploying Gemma 3 4B Vision-Language Model
         on Mobile Devices},
  author={AtomGradient},
  year={2026},
  url={https://github.com/AtomGradient/swift-gemma-cli}
}