chicagoypark/second-ckpts

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

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 (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).