MixFrag: Fragility-Guided Mixed-Precision Post-Training Quantization for Vision Transformers

Overview§

MixFrag is a fragility-guided mixed-precision post-training quantization (PTQ) framework for Vision Transformers (ViTs). Instead of applying uniform bit-widths to all layers, it estimates the sensitivity of each component (e.g., attention, MLP, normalization) to quantization using KL divergence between full-precision and quantized output distributions. Bit allocation is then formulated as a Multiple-Choice Knapsack Problem (MCKP) to meet a target bit budget while minimizing total fragility.

Methodology§

Fragility Estimation§

For each component $i$ (e.g., QKV projections, output projections, MLP blocks), MixFrag measures the KL divergence between the full-precision output distribution $P_i$ and the quantized output distribution $Q_i$ under a candidate bit-width $b$:

$$ \mathcal{F}_i(b) = \sum_{x} P_i(x) \log \frac{P_i(x)}{Q_i(x)} $$

This is computed over a small calibration set by quantizing each component in isolation while keeping others in full precision. The resulting fragility scores reflect how sensitive each component is to reduced precision.

Bit Allocation via MCKP§

Given a set of components $\mathcal{C}$ and candidate bit-widths $\mathcal{B}=\{2,4,6,8\}$, the goal is to minimize total fragility subject to a model-level bit budget $\mathcal{B}_{total}$:

$$ \min_{b_i \in \mathcal{B}} \sum_{i \in \mathcal{C}} \mathcal{F}_i(b_i) \quad \text{s.t.} \quad \sum_{i \in \mathcal{C}} b_i \cdot s_i \leq \mathcal{B}_{total} $$

where $s_i$ is the size (number of parameters) of component $i$. This is a standard Multiple-Choice Knapsack Problem, solvable via dynamic programming in $O(|\mathcal{C}| \times |\mathcal{B}| \times \mathcal{B}_{total})$.

Implementation Details§

• Calibration: 1,024 images from ImageNet-1K training set.
• Quantization scheme: Uniform affine quantization per-channel for weights and per-tensor for activations. Standard min–max or percentile-based clipping.
• Architectures tested: DeiT-S, DeiT-B, Swin-T, Swin-S, among others.
• Tasks: ImageNet classification, COCO object detection, and instance segmentation.

The following Python/PyTorch snippet illustrates the core fragility-guided allocation process:

import torch
import torch.nn as nn
from scipy.stats import entropy

def estimate_fragility(model, calib_loader, component_set, bits=[2,4,6,8]):
    fragility = {name: {b: 0.0 for b in bits} for name in component_set}
    for x, _ in calib_loader:
        # Full-precision outputs for all components
        full_outs = {name: get_output(model, x, name).detach() for name in component_set}
        for name in component_set:
            for b in bits:
                # Quantize only this component to b bits
                quantize_component(model, name, b)
                quant_outs = get_output(model, x, name).detach()
                # Compute KL divergence per sample and average
                frag = entropy(full_outs[name].flatten().cpu().numpy(),
                               quant_outs.flatten().cpu().numpy())
                fragility[name][b] += frag / len(calib_loader)
                restore_component(model, name)  # restore to full precision
    return fragility

def allocate_bits(fragility, component_sizes, target_budget, candidates=[2,4,6,8]):
    # DP for multiple-choice knapsack
    dp = {0: 0.0}
    choice = {}
    for name in fragility:
        new_dp = {}
        for cost, val in dp.items():
            for b in candidates:
                new_cost = cost + b * component_sizes[name]
                new_val = val + fragility[name][b]
                if new_cost <= target_budget and new_cost not in new_dp or \
                   new_dp.get(new_cost, float('inf')) > new_val:
                    new_dp[new_cost] = new_val
                    choice[(new_cost, name)] = (b, cost)
        dp = new_dp
    # backtrack to get bit allocations
    best_cost = min(dp)
    allocated = {}
    for name in reversed(fragility):
        b, prev_cost = choice[(best_cost, name)]
        allocated[name] = b
        best_cost = prev_cost
    return allocated

Results§

MixFrag achieves competitive ImageNet-1K accuracy under mixed-precision budgets, and state-of-the-art results on COCO detection/segmentation. Notably, under the challenging MP3/MP3 setting (3-bit weights and activations with mixed precision), it improves the previous best method by up to 9.6 AP. The learned bit allocations correlate strongly with the fragility scores, confirming that the KL-divergence metric is a reliable proxy for layer-wise sensitivity.

Key Takeaways§

• Component-level fragility quantification provides a principled way to decide bit-widths in PTQ.
• MCKP formulation ensures a globally optimal allocation under hard bit-budget constraints.
• The method is architecture-agnostic and works for both classification and dense prediction tasks.

MixFrag is a practical and effective approach for deploying ViTs on resource-constrained hardware, requiring only a small calibration set and no end-to-end training or gradient computation.