| --- |
| 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("vae/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). |
|
|
