Update compact quality model + graphs + report

a25163e verified 3 months ago

13.2 kB

Distilling an operator-video quality stack into a ~2 MB on-device model

This is the design write-up for video_benchmark.distill — how a heavy, multi-model video-quality scoring stack is compressed into a single ~2.5 M-parameter network that reproduces it in one forward pass and ships at ~2 MB on device (candle / Rust and Apple MLX), while staying honest about what it does and does not learn.

1. The problem

The production pipeline scores first-person / operator video frames (headband cameras, robotics teleop) on several quality axes. The accurate version runs a stack of models per frame:

Classical OpenCV metrics — brightness, sharpness, blur, exposure anomalies.
A 3-paradigm learned-IQA ensemble (via pyiqa): TOPIQ (CNN), MUSIQ (ViT) and CLIP-IQA+ (CLIP) — three different inductive biases voting on perceptual quality.
A MobileCLIP zero-shot scene classifier — is this a usable operator scene?

That is great for an offline batch, but it is several hundred milliseconds per frame and many hundreds of MB of weights — a non-starter for on-device / real-time use.

Goal: one compact model that emits all eight signals at once, fast, small enough to embed, and faithful enough to trust.

2. The teacher → student design

                 ┌─────────────── teacher stack (label generator) ───────────────┐
   frame  ─┬───► OpenCV: brightness · sharpness · blur · anomaly                  │
           ├───► pyiqa ensemble: TOPIQ(CNN) + MUSIQ(ViT) + CLIP-IQA+(CLIP)        │
           └───► MobileCLIP zero-shot scene usability                            │
                 └──────────────────────────► 8 per-frame targets (0..100) ──────┘
                                                          │  (distillation labels)
   frame  ──► frozen tiny backbone ──► embedding ─┐       ▼
                                                  ├──► fuse ──► residual-MLP trunk ──► 8 MLP heads ──► 8 scores
   frame  ──► 10 classical descriptors ───────────┘                                   (100·σ, bounded 0..100)

The student is CompactQualityNet (model.py):

Frozen backbone — a tiny preset MobileNetV3-Small (≈2.45 M params). Frozen so the embedding can be precomputed once and every training epoch is a cheap matrix op on the cache (data.py). Presets: tiny (default), micro (LCNet-050, sub-2 MB at every precision), clip (the legacy MobileCLIP-S0, max fidelity).
Descriptor fusion (see §4) — the embedding is concatenated with 10 cheap classical descriptors before the trunk.
Residual-MLP trunk — pre-norm residual blocks (x + W₂·GELU(W₁·LN(x))).
Per-target heads — one small 2-layer MLP per signal, so each stays independently calibrated while sharing the trunk.
Bounded output — 100·sigmoid(logit): scores start neutral at 50 and can never go negative or saturate the clamp, which keeps them calibrated.

The architecture (trunk dim, blocks, head dim, extra_dim) is persisted in the checkpoint, so infer.py rebuilds the exact network even if defaults later change.

3. Training (`train.py`)

A real mini-batch loop, not a single least-squares step:

kornia augmentation (augment.py) — each sampled frame is expanded into several views: a log-uniform Gaussian-blur sweep (perceptual blur is multiplicative, so log-uniform spreads samples evenly across mild→heavy), a mild motion blur (camera shake), plus colour/gamma jitter and sensor noise. View 0 is the untouched original. This multiplies the data and injects spread into otherwise near-flat signals so their distilled correlations become meaningful instead of n/a. Empirically, this lifts the rank-correlation metrics; aggressive degradations (strong motion blur, JPEG blocking) were measured to regress fidelity on this small corpus and are left out.
SGDR — CosineAnnealingWarmRestarts: the LR periodically "resurges" to escape plateaus.
EMA — an exponential-moving-average shadow of the trainable weights for a smoother final model (LayerNorm-only trunk, so no BatchNorm running-stat hazard).
Loss — SmoothL1(β=0.1) (β=1.0 degenerates to plain MSE on [0,1] targets) plus a differentiable Pearson-correlation term, so we optimise the exact fidelity metric (PLCC) the evaluation reports. Per-target weighting counts the deep signals double.
Ship-by-deep-PLCC — on tiny data the EMA is not guaranteed to win, so we keep the best EMA and the best raw weights and ship whichever scores higher held-out deep-PLCC.
accelerate — the same loop is correct on CPU / Apple MPS / CUDA.

4. The key idea: classical-descriptor fusion (`descriptors.py`)

A frozen ImageNet backbone is trained to be invariant to exactly what some quality signals measure: it normalises away absolute exposure, global contrast and colour cast. So brightness and exposure anomaly are nearly unrecoverable from its embedding — their distilled correlation collapsed to ~0.

Fix: a hybrid hand-crafted + deep representation. We concatenate a fixed 14-D vector of normalised classical statistics to the embedding:

Photometric / focus (10): luma mean/std, dark/bright fraction, RMS contrast, colourfulness (Hasler–Süsstrunk), saturation, Laplacian-variance sharpness, Canny edge density, Tenengrad.
BRISQUE natural-scene statistics (4): on the MSCN field I = (luma − μ_local) / (σ_local + 1), a pristine image's coefficients are ~unit-Gaussian; blur/noise/compression push the variance, excess kurtosis and the horizontal/vertical neighbour-product means off their natural values — a strong, classic no-reference quality cue (Mittal et al., 2012).

All are computed identically at cache-build and inference time with fixed scales (no data-dependent statistics), so the model stays self-contained and reproducible.

Two design choices make this strictly beneficial (model._FusedInput):

Split projection. The embedding keeps its own projection; the descriptors enter through a separate projection that is added on. This keeps the big embedding matrix's row length block-aligned, so it int4/int8 block-quantises in GGUF/MLX (a single fat Linear over the concatenation would have an unaligned row and fall back to fp16 — ~4× larger; this is exactly the regression that pushed an early build to 2.38 MB).
Zero-initialised. The descriptor projection starts at zero, so fusion is an exact no-op at init and can only help — descriptors are used only insofar as they reduce the loss, never diluting the embedding the deep signals rely on.

Result: best composite scores of any build, clipiqa/anomaly/brightness lifted, deep-signal PLCC held — at negligible size cost.

5. Evaluation (`evaluate.py`)

IQA-grade reporting, with rigor appropriate to a small held-out set:

The standard quad per signal — PLCC (Pearson), SRCC (Spearman), KRCC (Kendall), plus MAE and RMSE (0–100 units).
VQEG logistic-fitted PLCC / RMSE — the IQA convention (ITU-T P.1401 / VQEG) maps predictions through a monotonic 5-parameter logistic before correlating, so a model isn't penalised for an arbitrary nonlinear score scale when its ranking is right. Reported next to the raw numbers.
Composite — agreement on the overall quality verdict (mean over signals that genuinely vary), the bottom-line number.
Deep-signal aggregates — mean over the signals a learned model actually earns (iqa, musiq, clipiqa, scene).
BCa bootstrap 95 % CI on the composite PLCC — a single estimate on ~190 frames hides real uncertainty. We use the bias-corrected-and-accelerated bootstrap (Efron, 1987), which corrects the percentile interval for the median-bias and skew a correlation's sampling distribution has near ±1; it degrades to a percentile bootstrap if BCa can't be formed.
Robustness sweep — blur frames at increasing σ and check the student tracks the teacher's degradation (not just degrades arbitrarily), reported as a tracking MAE.
Fidelity vs σ — deep-PLCC bucketed by blur severity: a clean frame is easy, heavy blur is the real test.
Honesty guards — a near-flat teacher signal (std < 5 on 0–100) is reported as n/a, never a fake correlation. The train/val split is leakage-free by clip.

All of this is rendered to report.md / report.txt and published to the Hub in the modern .eval_results/ format (PLCC/SRCC/KRCC/MAE/RMSE per signal + composites), keyed to a companion benchmark dataset.

6. Results (latest run)

Final tiny + fusion run on the operator-video corpus (see report.md for the full per-signal table with KRCC/RMSE and the blur sweep):

metric	value
Composite PLCC	0.899 (95 % BCa CI [0.854, 0.925])
Composite PLCC — VQEG logistic-fitted	0.902
Composite SRCC / KRCC	0.896 / 0.714
Deep-signal PLCC	0.69
`iqa` / `musiq` PLCC	0.88 / 0.84
`sharpness` / `blur` PLCC	0.90 / 0.86
Throughput	~120× the teacher stack
Params / int4 weights	2.45 M / ~1.9 MB
int4 GGUF / MLX file	2.00 MB / 2.40 MB

7. Deployment (`quantize.py`, `export_hf.py`)

Sub-fp16 packing with real, scale-aware size accounting (it models the fp16 fallback for tensors GGUF can't block-quantise, so the reported int4/int8 numbers match what the files actually weigh — not an optimistic params × bytes guess):

GGUF (candle / Rust) — ggml Q8_0 / Q4_0 + fp16 fallback. Arch + target order + descriptor names ride along as metadata so a candle program rebuilds the forward pass. int4 GGUF ≈ 2.0 MB.
MLX (Apple Silicon) — grouped affine quant, int8 + int4, at the empirically smallest group size per bit-width (smaller groups qualify more matrices and cut fp16 fallbacks). Verified with an MLX-native quantise→dequantise round-trip self-test.
safetensors / .pt — the canonical fp16 weights infer.py loads.
LICENSE / NOTICE — backbone-aware: MIT for the trained trunk + heads; the timm MobileNetV3 backbone is Apache-2.0 (apple-amlr only when a MobileCLIP tower is actually bundled); teacher models (BSD pyiqa, apple-amlr MobileCLIP) are used only to generate labels and are not redistributed.

8. Reproduce

uv sync --group distill                 # gguf + mlx export backends (optional)
uv run python -m video_benchmark.distill --videos videos --epochs 300 \
    --preset tiny --quantize int4 --gguf --mlx --fusion
uv run python -m video_benchmark.distill.infer frame.jpg   # run the distilled model
uv run pytest tests/test_distill.py tests/test_distill_model.py

9. Honest limitations

scene is flat on this corpus (only a handful of clips, all valid operator scenes), so its distilled correlation is n/a — a data limitation, not a model one. More diverse footage would give it real spread.
brightness / anomaly are exact OpenCV stats in production; the distilled head is only a unified read-out for them. Fusion lifts them but they are better computed directly.
The Hub benchmark uploads correctly, but HF only renders a live leaderboard once it allow-lists the eval.yaml — and no evaluation_framework enum value fits a custom distillation-fidelity eval, so that field is a documented placeholder.

10. References

The design borrows established methods rather than inventing them:

Knowledge distillation — Hinton, Vinyals & Dean, Distilling the Knowledge in a Neural Network, 2015.
BRISQUE / MSCN natural-scene statistics — Mittal, Moorthy & Bovik, No-Reference Image Quality Assessment in the Spatial Domain, IEEE TIP 2012.
VQEG logistic fit for IQA — VQEG, Final Report on the Validation of Objective Models of Video Quality Assessment, 2003; ITU-T Rec. P.1401.
BCa bootstrap — Efron, Better Bootstrap Confidence Intervals, JASA 1987.
SGDR warm restarts — Loshchilov & Hutter, SGDR: Stochastic Gradient Descent with Warm Restarts, ICLR 2017.
Weight averaging (EMA / Polyak) — Polyak & Juditsky, 1992; Izmailov et al., SWA,
Backbone & teachers — MobileNetV3 (Howard et al., 2019); MobileCLIP (Vasu et al., CVPR 2024); the learned-IQA ensemble TOPIQ (Chen et al., 2024), MUSIQ (Ke et al., ICCV 2021), CLIP-IQA (Wang et al., AAAI 2023), via pyiqa.
Colourfulness — Hasler & Süsstrunk, Measuring Colourfulness in Natural Images,
1. Tenengrad focus — Krotkov, 1987.
kornia — Riba et al., Kornia: an Open Source Differentiable Computer Vision Library for PyTorch, WACV 2020.
GGUF / ggml k-quants (llama.cpp) and MLX (Apple) for the on-device exports.