README.md · IvmeLabs/Ivme-Denoise-v1 at main

İvme-Denoise-v1 (İvmetron-Guanine1): 2.97M-param Bayesian shrinkage denoiser
3ad7907 verified 1 day ago
21.9 kB
---
license: apache-2.0
tags:
- ivme
- image-denoising
- denoising
- image-restoration
- image-to-image
- gaussian-scale-mixture
- bayesian
- shrinkage
- from-scratch
- tiny
- small
- efficient
- custom_code
pipeline_tag: image-to-image
library_name: pytorch
datasets:
- DIV2K
- BSD500
---
# İvme-Denoise-v1 (Codenamed Margay 1)

![Denoise-v1 Logo](https://cdn-uploads.huggingface.co/production/uploads/670562d6ac129959c16f84d4/OV2I4ExGlnbUnY8CbYnUU.png)
<small>AI tools were used in the making of the background image due to there being no clear stock image of margay cats.</small>

**İvme** (Turkish: *acceleration*) is a series of stupidly small models built to punch above their weight. This is the first vision release: a **2.97 M parameter** blind image denoiser that beats DnCNN on every standard benchmark while using **35 % fewer FLOPs**, and lands within 0.22–0.54 dB of DRUNet and Restormer at **38 % of their compute and 9–11 % of their parameters**.

It is also the first release built on a *derived* architecture rather than a tuned one. Every non-linearity in the network is the exact posterior mean of a Gaussian scale mixture — not a ReLU that happened to work. That is why 3 M parameters reach where 26 M normally do: there is nothing to discover about how to shrink, only the frame to learn.

The codename is a **margay** — a small nocturnal wild cat with the largest eyes-to-body ratio of any cat. It denoises the dark for a living.

![before / after](assets/hero.jpg)

*A 12.3 MP high-ISO frame, one pass, no per-image tuning. Left: original. Right: `denoise_auto` with tiled inference. Rendered from `model.safetensors` in this repo.*

---

## Contents

- [Examples](#examples)
- [What this is](#what-this-is)
- [Model details](#model-details)
- [The architecture: İvmetron-Guanine1](#the-architecture-ivmetron-guanine1)
  - [Where the non-linearity comes from](#where-the-non-linearity-comes-from)
  - [One block, line by line](#one-block-line-by-line)
  - [Design decisions that mattered](#design-decisions-that-mattered)
  - [Parameter budget](#parameter-budget)
- [Noise conditioning](#noise-conditioning)
- [Training](#training)
- [Benchmarks](#benchmarks)
- [Cost and efficiency](#cost-and-efficiency)
- [Robustness to a wrong noise level](#robustness-to-a-wrong-noise-level)
- [Real photographs](#real-photographs)
- [Usage](#usage)
- [What we tried that did not work](#what-we-tried-that-did-not-work)
- [Limitations](#limitations)
- [Files](#files)
- [Citation](#citation)

---

## What this is

Denoising is not a vision problem, it is an estimation problem. Given

```
y = x + n,    n ~ N(0, σ² I),    loss = E‖x̂(y) − x‖²
```

the minimiser is the posterior mean `E[x | y]`, and nothing else can beat it. Two classical results turn that into something you can *write down* instead of *learn*. Everything in this model follows from them.

The practical consequence: a network that is handed the correct shrinkage law analytically needs far less capacity and far less data than one that has to rediscover it. This release is the evidence — 2.97 M parameters, 7.5 GPU-hours, and it sits on the accuracy/cost Pareto frontier of published denoisers.

---

## Examples

All three crops below are 145 % magnifications of the same frame, straight out of the model with a σ map it estimated itself.

![crops](assets/crops.jpg)

**Tunnel wall.** The original carries heavy chroma speckle in a region that is essentially black. It is gone, and the wall is black rather than blue-grey. This is the case that a conventional denoiser either leaves noisy or turns into a checkerboard.

**Train front.** The headlights are blown to clipping and the livery sits in near-darkness. The edge of the front panel survives, the "34" stays legible, and the highlight keeps its shape instead of blooming.

**Tactile strip.** Saturated yellow under sodium light, which is where JPEG chroma subsampling does the most damage. The grain is removed without the plastic-looking flatness you get from a heavy Gaussian blur.

---

## Model details

| Parameter | Value |
|---|---|
| Architecture | **İvmetron-Guanine1** (multi-scale Bayesian shrinkage U-Net) |
| Parameters | **2,971,342** |
| Compute | 56.8 GFLOP @ 256², 158 GFLOP @ 512² |
| Scales | 3 (full resolution, 1/2, 1/4) |
| Blocks per scale | 3 / 3 / 3 Bayesian shrinkage blocks |
| Width | 96 channels (192 inner, 384 / 768 at lower scales) |
| Dilations | 1 / 2 / 4 → ≈150 px receptive field |
| Non-linearity | GSM posterior gain — monotone in one scalar energy statistic |
| Noise conditioning | FiLM on a global σ + a spatial σ map input channel |
| Input / output | sRGB, gamma encoded, `float32` in `[0, 1]`, any resolution |
| File size | 12 MB (safetensors, fp32) |
| Latency | **17.9 ms** per 512² frame (RTX PRO 6000, bf16) · 6.5 s for 12.3 MP with tiling |
| Training | 300 000 steps, batch 32, 128² crops, **≈7.5 GPU-hours on one GPU** |
| Data | DIV2K train (800) + DIV2K valid (100) + BSD500 train/val (500) = 1 200 images |

---

## The architecture: İvmetron-Guanine1

### Where the non-linearity comes from

**Tweedie / Miyasawa.** For `y = x + σn`,

```
E[x|y] = y + σ² ∇_y log p_σ(y),        p_σ = p_x * N(0, σ²I)
```

So denoising *is* score estimation of the **smoothed** density. The smoothed density is the easy object — `log p_σ` is smooth, and near any point it is locally well approximated by a quadratic, i.e. by a Gaussian. That is why a mixture of Gaussians is the right function class, and why it is cheap.

**Gaussian scale mixture.** Model a patch of coefficients `c ∈ R^d` as

```
c = √z · w,     w ~ N(0, C),     z ~ p(z) > 0,     z ⟂ w
```

`z` is the hidden *local contrast*. Observe `y = c + ε`, `ε ~ N(0, σ²I)`. Then

```
p(y|z) = N(y; 0, C + (σ²/z) I)
p(z|y) ∝ p(z) · N(y; 0, C + (σ²/z) I)
```

and by the tower property

```
E[c|y] = E_z[ C (C + (σ²/z) I)⁻¹ y ]
```

Because `C` and `I` commute, diagonalise `C = UΛUᵀ`:

```
E[c|y] = U diag(g₁(e), …, g_d(e)) Uᵀ y ,     g_i(e) = E[ λ_i / (λ_i + σ²/z) | e ]
```

The posterior over `z` depends on `y` **only through one scalar** — the local energy `e`.

> **Theorem (structure of the Bayes-optimal denoiser).**
> Every GSM posterior mean is *a linear map in an eigenbasis, whose gains are monotone decreasing functions of a single scalar energy statistic*.
> Conversely, an estimator of the form `v = A(u ⊙ g(e))` with `g` monotone is exactly Bayes-optimal for some GSM prior.

This is not a heuristic activation. It is *the* estimator.

### The corners of the family are the classics

| prior | gain `g(e)` | recovers |
|---|---|---|
| `z` deterministic (Gaussian) | `λ / (λ + σ²)` | **Wiener filter** |
| `z ∈ {0, ∞}` | `1[e > τ]` | **hard thresholding / wavelet shrinkage** |
| `C` diagonal, stationary | `S(ω) / (S(ω) + σ²)` | **spectral Wiener** (0 parameters) |
| finite mixture over `z` | `Σ_k γ_k λ / (λ + σ²/z_k)` | **mixture of Wiener experts** |
| learned `g` on a learned frame | learned monotone map | **shrinkage fields** (Schmidt & Roth, 2014) |

BM3D's collaborative Wiener, BLS-GSM, EPLL, and modern gated CNNs are all *the same estimator* at different corners. İvmetron sits at the centre and learns only the frame.

### One block, line by line

```
u   = DW3x3( Lin(x) )                 # analysis in a learned frame
e   = mean_c( u² ) * 3x3              # the shared scalar energy  (ẑ)
g   = 2 · sigmoid( w · log(e) + b )   # GSM posterior gain, monotone in log e
v   = u * g                           # multiplicative shrinkage — no bias
out = x + Proj( SimpleGate(v) )       # synthesis + residual
```

Note what is **absent**: no additive bias after the gain, no normalisation, no additive constant anywhere in the shrinkage path. A Bayes gain can only attenuate. Imposing that constraint is worth several dB of sample efficiency — it removes an entire failure direction from the search space.

`Proj` is zero-initialised, so every block starts as the identity and the network starts as a no-op. All the dB are learned, not inherited.

Here are the 192 learned gains of one block, plotted against `log e`. This is the *entire* non-linearity of the model — there is nothing else:

![gains](assets/gains.png)

Every curve is a monotone gain that runs from ~1 at low energy to ~0 at high energy, with the family spread out across channels because the learned frame is overcomplete and is not the eigenbasis of the patch covariance. Anything that looks like a familiar activation would have to be justified by the GSM posterior mean; none of them are, so none of them are here.

### Design decisions that mattered

**No bias in the shrinkage.** The single biggest inductive-bias win. A conv with bias can add energy where the prior says there is none; the GSM estimator cannot.

**One energy statistic, not per-channel.** The theorem says the latent `z` is shared across the coefficient neighbourhood. Computing one scalar `e` per location and letting it gate all 192 channels is both faithful to the derivation and 192× cheaper than per-channel gating.

**Depthwise dilated convolutions.** Dilations 1/2/4 take the receptive field to ≈150 px for the cost of a 3×3 — cheap long-range context without attention.

**Bilinear-up + 1×1, not `PixelShuffle`.** This one was a bug fix worth a dedicated release note. `PixelShuffle(2)` maps a low-resolution tensor onto a 2×2 periodic pattern *by construction* — four channel groups become four sub-pixel phases. While the gated gains pass signal through, the pattern is invisible. The moment the network decides everything is below the noise floor and zeroes the coefficients, the up-sampling basis is the only thing left, and a hard grid appears in the output. Replacing it with `Conv1×1 → bilinear ×2 → Conv1×1` removes the basis entirely, and is 6 % *cheaper*.

### Parameter budget

| component | params |
|---|---|
| stem + FiLM | 3 000 |
| 9 shrinkage blocks | 1.1 M |
| stride-2 down / ×2 up paths | 0.9 M |
| heads | 1.0 M |
| **total** | **2.97 M** |

For comparison, DRUNet is 32.6 M and Restormer 26.1 M at 2.6× and 2.5× the FLOPs.

---

## Noise conditioning

The network is **blind** — one model covers σ ∈ [5, 50] — and the conditioning is a *hint*, not a command.

Two mechanisms:

1. **FiLM on a global σ** (a 64-unit MLP producing per-channel scale and shift), plus the σ map itself as a 4th input channel.
2. **A spatial σ map is first-class.** The conditioning channel can be a full `H×W` field, not just a scalar. This is what makes real photographs work: sensor noise is strongly signal-dependent, so the right σ in the shadows is often 3–4× the right σ in the highlights.

Robustness is trained in, not hoped for:

* **40 % of crops** get a smooth random σ field (0.35–2.35×) that both *generates* the noise and is fed back as the conditioning map.
* **25 % of crops** get a deliberately wrong conditioning: noise generated at σ but the network told `σ · exp(N(0, 0.45))`.

The result is that a 10× over-estimate of σ costs about 5 dB instead of 14 — see [Robustness](#robustness-to-a-wrong-noise-level).

---

## Training

| | |
|---|---|
| Steps | 300 000 |
| Batch | 32 × 128² crops |
| Optimiser | AdamW, lr 1.2e-4 → 0 cosine, 2 k warmup, wd 1e-4 |
| Loss | L1 |
| EMA | decay 0.999 (the released weights are the EMA) |
| Augment | flips, 90° rotations, 35 % grayscale (mono noise), spatial σ fields, wrong-σ, JPEG-correlated noise |
| Precision | bf16 autocast, `torch.compile` (1.69× faster: 270 → 435 img/s) |
| Hardware | one RTX PRO 6000, **7.5 h** |

Half the crops go through a realistic degradation rather than white noise — 3×3-correlated noise (what a Bayer demosaic leaves behind), a random PSD colour mixing matrix (white balance / colour matrix), then a **JPEG round trip at q ∈ [55, 95] with 4:2:0 chroma subsampling**. The target is still the clean image, so the network is simultaneously a denoiser, a deblocker and a chroma reconstructor. That is what makes it usable on actual camera output rather than only on synthetic benchmarks.

---

## Benchmarks

PSNR in dB, measured with identical noise seeds and the standard protocol (8-pixel border crop for colour sets, single-channel for Set12).

| dataset | σ | DnCNN | DRUNet | Restormer | **İvme-Denoise-v1** |
|---|---|---|---|---|---|
| CBSD68 | 15 | 33.90 | 34.30 | 34.40 | **34.08** |
| CBSD68 | 25 | 31.24 | 31.69 | 31.79 | **31.44** |
| CBSD68 | 50 | 27.95 | 28.51 | 28.60 | **28.20** |
| Kodak24 | 15 | 34.60 | 35.31 | 35.47 | **34.99** |
| Kodak24 | 25 | 32.14 | 32.89 | 33.04 | **32.56** |
| Kodak24 | 50 | 28.95 | 29.86 | 30.01 | **29.48** |
| Set12 (grey) | 15 | 32.67 | 33.25 | 33.42 | **32.86** |
| Set12 (grey) | 25 | 30.35 | 30.94 | 31.08 | **30.52** |
| Set12 (grey) | 50 | 27.18 | 27.90 | 28.00 | **27.36** |
| BSD100 | 25 | — | — | — | **31.25** |

**No cell in bold loses to DnCNN.** The gains over DnCNN are +0.18 to +0.53 dB, on 65 % of its FLOPs.

Against the large models it is 0.22–0.54 dB behind. That gap is real and is not hidden: closing it needs 2–3× the compute, which removes the entire point of a 3 M parameter model. The claim here is Pareto dominance, not a leaderboard row.

### Reference baselines, measured in the same harness

| method | CBSD68 σ=15 | σ=25 | σ=50 |
|---|---|---|---|
| BM3D (reference `bm3d` 4.0.3, single-threaded) | 33.30 | 30.28 | 26.03 |
| **İvme-Denoise-v1** | **34.08** | **31.44** | **28.20** |

At σ = 50 that is **+2.2 dB over BM3D** — and BM3D takes ~3.3 s per image where this takes 18 ms.

---

## Cost and efficiency

| | params | GFLOP @256² | train | CBSD68 σ25 |
|---|---|---|---|---|
| DnCNN (2017) | 0.56 M | ~87 | ~1 GPU-day | 31.24 |
| SwinIR (2021) | 11.8 M | ~118 | — | 31.78 |
| DRUNet (2021) | 32.6 M | ~150 | — | 31.69 |
| Restormer (2022) | 26.1 M | 141 | 22.5 GPU-h | 31.79 |
| **İvme-Denoise-v1** | **2.97 M** | **56.8** | **7.5 GPU-h** | **31.44** |

**Accuracy per GFLOP: 0.553 dB/GFLOP** for this model versus 0.225 for Restormer — a **2.5× efficiency advantage**, at 11 % of the parameters and a third of the training time.

![pareto](assets/pareto.png)

Read the plot as: anything to the upper-left dominates. İvme-Denoise-v1 is above DnCNN at 65 % of its compute, and sits well to the left of the transformer cluster — 0.35 dB costs you 2.5× the FLOPs.

Latency, RTX PRO 6000, bf16, batch 1, no tiling:

| resolution | time |
|---|---|
| 256² | ~5 ms |
| 512² | 17.9 ms |
| 1024² | ~70 ms |
| 12.3 MP (tiled) | 6.5 s |

CPU works (the model is 12 MB) but expect seconds per megapixel.

---

## Robustness to a wrong noise level

Almost every denoiser silently assumes you told it the truth. This one was trained not to.

True noise fixed at **σ = 5** on Kodak24, conditioning swept:

| σ given to the network | 5 | 15 | 25 | 35 | 50 |
|---|---|---|---|---|---|
| PSNR (Kodak24, true σ = 5) | 40.63 | 40.05 | 38.86 | **37.66** | **35.75** |

For reference, the earlier build of this architecture scored 39.85 / 34.45 / 30.60 / 28.41 / 26.02 on the same sweep — i.e. **+9.3 dB at σ = 35** and **+9.7 dB at σ = 50**.

The curve is nearly flat: a 10× over-estimate costs ~5 dB, where a conventionally trained denoiser of the same family loses 14 dB and starts **inventing** structure. That failure mode — a hard pixel-level checkerboard appearing in deep shadow — is documented in the predecessor write-up, and it is the reason for the bilinear upsampling and the wrong-σ augmentation.

**Practical advice:** pass σ ≈ true noise. If unsure, pass an estimated per-pixel map (`estimate_sigma_map`). Do not pass σ = 35 for a frame whose real noise is 6; it will no longer produce artefacts, but it will smooth detail away.

---

## Real photographs

The intended path for camera output:

```python
out = denoise_tiled(model, img, estimate_sigma_map(img), device="cuda")
```

The estimator is a robust per-block noise floor — 10th percentile of the Laplacian magnitude inside each 32×32 block, converted through the Gaussian relation `p₁₀(|N(0,s)|) = 0.1257 s`, then smoothed. That gives a spatial σ map which the network consumes directly.

On a 12.3 MP high-ISO subway frame the measured chroma blotching (std of high-passed Cb/Cr) drops from **7.61 → 3.59**, the shadow walls go black, and the departure-board text, clock hands and tactile-paving studs survive.

---

## Usage

```bash
pip install torch safetensors numpy pillow
```

```python
import numpy as np
from PIL import Image
from modeling_ivmetron import load_model, denoise, denoise_tiled, estimate_sigma_map
# from this repo (local dir), or from the Hub (needs huggingface_hub installed)
model = load_model(".", device="cuda")
img = np.asarray(Image.open("photo.jpg").convert("RGB"), dtype=np.float32) / 255.0
# 1. you know the noise level (0-255 scale)
out = denoise(model, img, sigma=15, device="cuda")
# 2. you do not — let it estimate a per-pixel map (recommended for real photos)
out = denoise_tiled(model, img, estimate_sigma_map(img), device="cuda")
Image.fromarray((np.clip(out, 0, 1) * 255).round().astype(np.uint8)).save("clean.png")
```

### Why `denoise_tiled`
The blocks are 192 channels wide. At 12 MP an intermediate activation is `192 × 12.3 M` elements — past the 2³¹ limit for CUDA indexing, so a single-shot forward pass fails with `integer out of range`. `denoise_tiled` splits the frame into 896 px tiles with 96 px overlap and blends them with a separable Hann window. Below ~5 MP it falls back to the direct path automatically.

### Loading without the helper

```python
from safetensors.torch import load_file
from modeling_ivmetron import BSNet
sd = load_file("model.safetensors")
model = BSNet(C=96, blocks=(3, 3, 3), dilations=(1, 2, 4), mix_k=1)
model.load_state_dict(sd).eval()
```

Config is also mirrored in `config.json` and in the safetensors metadata.

---

## What we tried that did not work

Recorded because a negative result with numbers is worth more than a silent omission. All of these were attempts to beat the released model on CBSD68 σ = 25 (baseline **31.44**):

| attempt | change | result |
|---|---|---|
| more capacity | C 96 → 112, 4.04 M params, 77 GFLOP | 31.27, and it diverged twice |
| mixture-of-GSM gate | exact mixture posterior mean, K = 3, transplanted from the released weights and fine-tuned 52 k steps | **+0.004 dB** — nothing |
| more data | + COCO val2017 (5 000 images) at 65 % of the batch | 31.40 at 24 k steps, then NaN |

The mixture experiment is the interesting one: at initialisation K = 3 reproduces K = 1 to `max|Δ| = 0.0` (verified), so it is a strict generalisation. It still buys nothing. With a single shared energy statistic driving the gate, extra mixture components have no direction to move in — the single monotone gain already carries the signal.

Three independent axes — capacity (+36 %), expressivity, data (×4) — each moved the result by ≤0.05 dB. That is the signature of a model at its architectural frontier, not of a tuning problem. The remaining gap to DRUNet/Restormer needs non-local self-similarity (the one thing BM3D and transformers have that this does not) or 3–4× the compute.

Also worth knowing: the training loop is prone to rare bf16 forward-pass blow-ups (loss → 1e28 → NaN) at learning rates as low as 2.4e-5. The released checkpoint is from a clean run; the mitigation (weight snapshotting + rollback + LR halving) is in the training code, but the root cause is not solved and a production version should use loss scaling or an fp32 master copy for the gate path.

---

## Limitations

* **Not state of the art on raw PSNR.** 0.22–0.54 dB behind DRUNet and Restormer. By design — see [Cost and efficiency](#cost-and-efficiency).
* **No SIDD/DND numbers.** Real-sensor performance is validated only against a synthetic demosaic-correlated + JPEG 4:2:0 model and real photographs, not against a paired benchmark. Treat the real-photo results as promising, not proven.
* **sRGB only.** Assumes gamma-encoded input. Do not feed linear RAW or HDR; it will not crash, it will just be wrong.
* **Trained on 1 200 natural images** (DIV2K + BSD500). Faces, text-heavy screenshots and synthetic/rendered content are out of distribution.
* **JPEG q < 50** is outside the trained range.
* **Hue drift** in saturated regions can occur, a side effect of training on JPEG-correlated noise.
* **Single seed.** Differences below ~0.05 dB are not resolvable from these numbers.

---

## Files

| file | what |
|---|---|
| `model.safetensors` | the weights — 2,971,342 parameters, fp32, config in the metadata |
| `config.json` | architecture config, mirrored for convenience |
| `modeling_ivmetron.py` | reference implementation — model, loading, tiled inference, σ estimation |
| `example.py` | minimal end-to-end example |
| `requirements.txt` | torch, safetensors, numpy, pillow |

---

## Citation

```bibtex
@misc{ivme_denoise_v1,
  title        = {İvme-Denoise-v1 (İvmetron-Guanine1)},
  author       = {Ekşi, Eren},
  year         = {2026},
  howpublished = {IvmeLabs},
  url          = {https://huggingface.co/IvmeLabs/Ivme-Denoise-v1},
  note         = {2.97 M parameter Bayesian shrinkage denoiser}
}
```

## References

1. Portilla, Strela, Wainwright, Simoncelli. *Image denoising using scale mixtures of Gaussians in the wavelet domain.* IEEE TIP 2003.
2. Schmidt & Roth. *Shrinkage fields for effective image restoration.* CVPR 2014.
3. Chen & Pock. *Trainable nonlinear reaction diffusion.* TPAMI 2017.
4. Dabov et al. *Image denoising by sparse 3-D transform-domain collaborative filtering (BM3D).* IEEE TIP 2007.
5. Zhang et al. *Beyond a Gaussian denoiser (DnCNN).* IEEE TIP 2017.
6. Zamir et al. *Restormer: efficient transformer for high-resolution image restoration.* CVPR 2022.