---
license: apache-2.0
tags:
  - vae
  - autoencoder-kl
  - medical-imaging
  - latent-diffusion
library_name: pytorch
---

# kl16_mm3ch_gan_balanced — 130k

3-channel KL-16 VAE (LPIPS + patch-GAN) for multimodal medical imaging, trained with **balanced
cohort sampling** across BraTS 2023, SynthRAD 2023, CHAOS and AMOS 22.

Continued from `kl16_mm3ch_gan_410k.ckpt` (weights only), which had been trained with
*proportional* sampling and is therefore heavily BraTS-dominated. This checkpoint re-balances the
cohort marginal to 25% each. Snapshot at **optimizer step 130,000**.

## Architecture

`AutoencoderKL` — the LDM/CompVis encoder–decoder, matched to the MAR `kl16` layout:

| | |
|---|---|
| `embed_dim` / `z_channels` | 16 / 16 (`double_z=True`) |
| `ch` / `ch_mult` | 128 / `[1, 1, 2, 2, 4]` |
| `num_res_blocks` | 2 |
| `attn_resolutions` | `[16]` in the **encoder**, `[]` in the **decoder** |
| resolution / channels | 256 × 256, `in_channels=3`, `out_ch=3` |
| parameters | 66.46 M (312 tensors, fp32) |

Downsampling factor f = 16, so a 256×256×3 input maps to a 16×16×16 latent.

> The asymmetric attention (encoder has it, decoder does not) is inherited from the MAR `kl16`
> checkpoint this lineage started from. Instantiating a stock LDM `AutoencoderKL` with attention in
> the decoder will produce missing/unexpected keys.

## Usage

```python
import torch
from ldm.modules.diffusionmodules.model import Encoder, Decoder

DDCONFIG = dict(double_z=True, z_channels=16, resolution=256, in_channels=3, out_ch=3,
                ch=128, ch_mult=[1, 1, 2, 2, 4], num_res_blocks=2,
                attn_resolutions=[16], dropout=0.0)

encoder = Encoder(**DDCONFIG)
decoder = Decoder(**{**DDCONFIG, "attn_resolutions": []})   # no attention in the decoder

sd = torch.load("kl16_mm3ch_gan_balanced_130k.ckpt", map_location="cpu")["model"]
# keys: encoder.* / decoder.* / quant_conv.* / post_quant_conv.*
```

Inputs are 3-channel, 256×256, scaled to `[-1, 1]`. Single-channel medical slices are replicated to
3 channels to match the backbone and the LPIPS/VGG perceptual loss.

## Training data

`content=both` — anatomical images **and** segmentation masks are both treated as samples.

| cohort | subjects (train) | train slices | natural share | sampled share |
|---|--:|--:|--:|--:|
| BraTS 2023 (GLI) | 1,188 | 736,560 | 87.5% | 25% |
| SynthRAD 2023 (Task 1) | 356 | 61,333 | 7.3% | 25% |
| AMOS 22 (abdomen) | 306 | 39,979 | 4.7% | 25% |
| CHAOS (abdomen) | 34 | 4,242 | 0.5% | 25% |
| **pooled** | **1,884** | **842,114** | | |

BraTS contributes 5 streams per subject (`t1n`, `t1c`, `t2w`, `t2f`, `seg`); SynthRAD contributes
MR + CT; CHAOS and AMOS contribute CT/MR plus liver masks. Sampling is multinomial with
replacement at a fixed per-cohort marginal, so CHAOS is oversampled roughly 25× per epoch and
BraTS is undersampled to ~0.14×.

Preprocessing: MRI normalized by per-volume non-zero (1, 99) percentiles, CT by a fixed
`[-1000, 1000]` HU window; body-bounding-box crop, pad to square, resize to 256, output in
`[-1, 1]`.

## Training configuration

| | |
|---|---|
| initialization | `kl16_mm3ch_gan_410k.ckpt`, weights only (step counter restarts at 0) |
| loss | L1 + LPIPS(VGG16) + KL (`kl_weight 1e-6`) + hinge GAN |
| discriminator | 3-layer patch critic, `disc_weight 0.5`, adaptive `d_weight = ‖∇nll‖ / ‖∇g‖ × 0.5` |
| optimizer | Adam(β = 0.5, 0.9), **constant lr 2.88e-4**, no schedule, no warm-up, fp32 |
| effective batch | 96 |
| step | 130,000 |

The 410k checkpoint carried no discriminator weights, so the critic was re-initialized. It was
first warmed up for 2,000 steps with `disc_weight = 0` — the critic trains while the generator
receives exactly zero adversarial gradient — before the full GAN objective was enabled.

## Caveats

- **The latent scale must be re-fit.** The scale factor `0.685155` belongs to
  `kl16_mm3ch_gan_410k.ckpt` and is **not** valid here. Re-estimate the latent std before pairing
  this VAE with any latent-space model.
- **Reconstruction-only validation drifts upward under adversarial training.** Validation scores
  pure L1 + LPIPS with no GAN term; it moved from 0.0211 (10k) to ~0.0222 and then flattened. That
  is the expected fidelity-for-sharpness trade, not divergence. Judge this checkpoint perceptually
  or on a downstream task, not by that scalar.
- **Research use only.** Trained on public research datasets; not validated for, or intended for,
  clinical use.

## Provenance

Derived from [CompVis/latent-diffusion](https://github.com/CompVis/latent-diffusion)
(`AutoencoderKL`, `LPIPSWithDiscriminator`) and the MAR `kl16` VAE lineage. Source datasets are
subject to their own licenses and access terms (BraTS 2023, SynthRAD 2023, CHAOS, AMOS 22).