Architecture

The generator, the critics and the game between them

Every model in this study shares one generator — a 7.76 M-parameter 2D U-Net. What changes between runs is the network sitting opposite it: its kernel size, its stride, its downsampling strategy, and how much of a slice one verdict is allowed to cover.

Model A no critic Model B 3×3 / stride 2 Model C 3×3 / stride 1 + pool Model D 4×4 / stride 2

The generator

2D U-Net — encoder, bottleneck, decoder

Four pooling levels, doubling filters on the way down and halving on the way up, with a skip connection at every level so the decoder can recover the spatial detail that pooling destroyed.

Every convolution uses padding='same', ReLU and he_normal initialisation. Only the bottleneck carries dropout. Skip connections are channel-axis concatenations, so the decoder convolution that follows sees both the upsampled semantic features and the high-resolution encoder features.

Total parameters

7,759,908

~29.6 MB at float32

Pooling levels

128 → 64 → 32 → 16 → 8

Receptive field

~140px

wider than the 128 px input — the bottleneck sees the whole slice

Design decisions, and what they cost

Choice	Reason	Trade-off accepted
2D, not 3D	3D activations plus a discriminator exceed the available GPU memory	no inter-slice context
128×128 input	power of two divides evenly through four pooling levels	detail lost on the small necrotic core
2 channels, not 4	T1ce and FLAIR carry near-complementary information	T1 and T2 residual signal discarded
32→512 filters	half the width of the original U-Net; fits one patient per step	lower capacity than a full-width model
Dropout only at the bottleneck	regularise the widest layer without disturbing skip features	encoder can overfit
`UpSampling2D` + Conv	avoids the checkerboard artefacts of transposed convolution	no learned upsampling kernel

The full Keras plot_model graph of the U-Net, listing every layer with its input and output tensor shapes from the 128 by 128 by 2 input down to the 128 by 128 by 4 softmax output. — **The Keras layer graph** as emitted by `plot_model` during training — every layer with its exact tensor shapes. Scroll inside the frame, or click to open full size.

The adversarial objective

How a critic changes what the U-Net optimises

The generator is unchanged. What changes is the loss: alongside cross-entropy it now receives a gradient from a network that has been trained to recognise expert annotations.

Both networks are updated inside one @tf.function-compiled step using two GradientTape contexts, with separate Adam optimisers at lr 2×10^-4, β₁ = 0.5.

Why λ = 100

The generator's loss has two terms with very different scales. The adversarial term is a binary cross-entropy on a single sigmoid — order 1. The segmentation term is a categorical cross-entropy that, on a well-fit model, sits around 0.02. Left unweighted, the adversarial term would be ~50× louder and the generator would optimise for a critic-pleasing mask rather than a correct one.

λ = 100 restores the balance: cross-entropy stays the dominant signal and the adversarial term acts as a tie-breaker among pixel-equivalent solutions, nudging the generator toward the one that looks hand-drawn. This is the same weighting pix2pix uses for its L1 term, for the same reason.

@tf.function
def train_step(mri_images, real_masks):
    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
        fake_masks  = generator(mri_images, training=True)

        real_output = discriminator([mri_images, real_masks], training=True)
        fake_output = discriminator([mri_images, fake_masks], training=True)

        total_disc_loss = (bce_loss(tf.ones_like(real_output),  real_output)
                         + bce_loss(tf.zeros_like(fake_output), fake_output))

        gen_adv_loss = bce_loss(tf.ones_like(fake_output), fake_output)
        gen_seg_loss = cce_loss(real_masks, fake_masks)
        total_gen_loss = gen_adv_loss + (LAMBDA_SEG * gen_seg_loss)   # LAMBDA_SEG = 100

    generator_optimizer.apply_gradients(zip(
        gen_tape.gradient(total_gen_loss, generator.trainable_variables),
        generator.trainable_variables))
    discriminator_optimizer.apply_gradients(zip(
        disc_tape.gradient(total_disc_loss, discriminator.trainable_variables),
        discriminator.trainable_variables))

The variable under test

Three critics, three failure modes

This is the whole experiment. Same generator, same data, same epochs — only the critic changes. Two of the three configurations break, and both break for reasons that are visible in the architecture before a single epoch is run.

The geometry behind the paradox. A single highlighted patch on the right contains one flat colour. Whether that patch belongs to a perfect mask or one shifted 40 px off-target is not decidable from inside the patch — so a patch critic cannot supply the global shape prior the study was trying to add.

Layer stacks

Model D · 4×4, stride 2 champion

Concat(MRI, mask)      128²×6
Conv 4×4 s2   64      64²
LeakyReLU 0.2
Conv 4×4 s2  128      32²
LeakyReLU 0.2
Conv 4×4 s2  256      16²
LeakyReLU 0.2
Flatten                65,536
Dropout 0.3
Dense 1 sigmoid             1

727,489 parameters. 4 ÷ 2 = 2, so every interior column is seen by exactly two windows — symmetric feature extraction, no checkerboard. Dropout on the flattened vector keeps the critic from overpowering the generator early.

Model B · 3×3, stride 2 artefacts

Concat(MRI, mask)      128²×6
Conv 3×3 s2   64      64²
LeakyReLU 0.2
Conv 3×3 s2  128 +BN  32²
Conv 3×3 s2  256 +BN  16²
Conv 3×3 s1  512 +BN  16²
Conv 3×3 s1    1      16²
sigmoid  → verdict grid

1,560,897 parameters. Classic pix2pix shape with a smaller kernel — and 3 ÷ 2 = 1.5, so coverage alternates 2, 1, 2, 1 and the critic carries a period-2 spatial bias into every gradient.

Model C · 3×3, stride 1 the paradox

Concat(MRI, mask)      128²×6
Conv 3×3 s1   64      128²
MaxPool 2×2           64²
Conv 3×3 s1  128 +BN  64²
MaxPool 2×2           32²
Conv 3×3 s1  256 +BN  32²
MaxPool 2×2           16²
Conv 3×3 s1  512 +BN  16²
Conv 3×3 s1    1      16²

1,559,937 parameters (1,558,145 trainable). Bias-free convolutions, N(0, 0.02) init. Stride 1 restores uniform coverage — and costs ~4× the FLOPs, 211 min of training, and the study's worst global coherence.

The intended fix that backfired. Model C exists because Model B's 3×3/stride-2 kernel does not tile its input evenly. Stride 1 plus max-pooling removes that bias — but pooling also makes the critic locally translation-invariant, discarding exactly the positional information needed to judge whether a mask sits in the right place. Its discriminator loss climbed from 9.06 to 12.48 over 35 epochs instead of falling, and its SSIM became the worst in the study even as its local segmentation loss became the best. That is the PatchGAN Paradox. See the divergence in the training curves →

Keras layer graph of the 4 by 4 global discriminator, showing the two inputs concatenated, three strided convolutions with LeakyReLU, a flatten, dropout, and a single sigmoid dense output. — **Model D’s critic, as Keras built it.** Two named inputs (`mri_image`, `segmentation_mask`) join at the concatenate layer, pass through three 4×4 stride-2 blocks and reduce to a single verdict. Note the implemented head is a `Flatten` → `Dense(1)` global verdict rather than a patch grid; the kernel/stride argument applies to the convolutional stack either way.

Reproducibility

Complete training configuration

Setting	Model A (U-Net)	Models B / C / D (cGAN)
Input	128×128×2	128×128×2
Output	128×128×4 softmax	128×128×4 softmax
Generator params	7,759,908	7,759,908
Loss	categorical cross-entropy	BCE + 100 · CCE (G), BCE (D)
Optimiser	Adam	Adam ×2 (separate)
Learning rate	1×10^-3	2×10^-4
β₁	0.9 (default)	0.5
LR schedule	ReduceLROnPlateau, factor 0.2, patience 2	none
Epochs	35	35
Steps per epoch	250	250
Batch	1 patient = 100 slices	1 patient = 100 slices
Training loop	`model.fit()`	`tf.GradientTape`, custom
Checkpointing	best `val_loss`, weights only	generator weights every 5 epochs
Init	`he_normal`	`he_normal` (G), varies (D)

Environment: Kaggle GPU notebooks, TensorFlow with Keras 3, Python 3.10. Every value above is read directly from the notebook source — browse the runs to see it in context.