Auto-Context Model (stacked refinement with previous predictions)

Overview

Auto-context is a way to give a patch-based model access to context that lies outside its patch. Stage 1 is trained on the raw input. Its predictions for whole volumes are then stacked onto the input as an additional channel, and stage 2 is trained on that stacked input, and so on 1612.05362 §2.2 1612.05362 fig. 3. The idea comes from Tu and Bai, who used it for segmentation classifiers; the paper's authors moved it to GAN regressors 1612.05362 §2.2.

Only one card on wikillm covers this concept: Medical Image Synthesis with Context-Aware Generative Adversarial Networks. Everything below is what that paper states or reports. The original Tu and Bai work has no card here, so it is named but not summarized.

What the variant on this page shares with the general pattern:

• Full-volume context. The stage-k output is a whole predicted volume, assembled from overlapping patch predictions averaged where they overlap 1612.05362 §2 1612.05362 §2.2.
• Context is a channel. The previous estimate is concatenated to the original input, so the next network sees two channels instead of one 1612.05362 §2.2.
• Sequential training. Each stage is trained only after the previous one has produced predictions 1612.05362 §2.2 1612.05362 fig. 3.

Variants

Two-stage ACM on a 3D context-aware GAN

Each stage is a patch-based 3D fully convolutional generator trained as a GAN generator, with an adversarial loss, an L2 loss and a 3D gradient difference loss 1612.05362 §2.1 1612.05362 eq. 7. Stage 1 sees MRI patches only. Later stages see MRI plus the previous CT estimate 1612.05362 §2.2. Two ACM iterations are used on both datasets 1612.05362 §4.2. The Paper vs code section of the card notes that the released 3D code has no auto-context stage and only a single-channel input, so ACM has to be built from the text 1612.05362 code: 3dganversion/g_model.py#L56.

The paper reports that both brain and pelvis results improve with more ACM iterations, mostly from the first to the second. The numbers are only plotted, not tabulated 1612.05362 §4.2 1612.05362 fig. 5. The reported numeric ablation is FCN versus GAN without ACM, not ACM on versus off 1612.05362 §4.1.

Variants not covered

The card does not describe other ACM designs, such as more than two stages, ACM in a non-GAN model, or stage outputs fed back in with a different combination rule. No other card on wikillm covers them, so they are not discussed here.

When to use which

Implementation sketch

Stage-wise training with a full-volume prediction between stages. The loop core is from the card's method description. Each stage's inputs are the original MRI plus the previous stage's output. Overlap-averaged inference follows the card's inference notes.

# Auto-context over stages  [arxiv:1612.05362 §2.2] [arxiv:1612.05362 fig. 3]
# Input:  paired mr_volumes, ct_volumes; num_stages (paper: 2)
# Output: one generator per stage

def train_auto_context(mr_volumes, ct_volumes, num_stages):
    generators = []
    context = None  # stage 1 has no context channel  [§2.2]
    for stage in range(num_stages):
        # patches: MRI (+ context) in, CT patch out  [§4]
        samples = extract_patches(mr_volumes, context, ct_volumes)
        in_channels = 1 if context is None else 2  # [§2.2]
        generator = train_gan(samples, in_channels)
        generators.append(generator)
        # full-volume predictions become the next stage's context  [§2]
        # must be run on every training subject before stage k+1  [§2.2]
        context = [predict_volume(generator, mr, c)
                   for mr, c in zip(mr_volumes, context or [None] * len(mr_volumes))]
    return generators

def predict_volume(generator, mr, context):
    # overlapping patches, averaged where they overlap  [§2]
    # stride is not stated in the paper  [code: 3dganversion/g_model.py#L178]
    output = zeros_like(mr)
    counts = zeros_like(mr)
    for centre in sliding_window(mr):
        inputs = stack_channels(mr, context, centre)
        output[centre] += generator(inputs)
        counts[centre] += 1
    return output / counts

Contradictions

None known. Only one card covers this concept. Within that card, the paper and the released code disagree: the paper describes ACM, and the code does not implement it 1612.05362 code: 3dganversion/g_model.py#L56.

Open questions the card records:

• Whether stage-k training context comes from a model that saw those same subjects, which could make the context optimistic, is not stated 1612.05362 §2.2.
• The number of stages beyond two, and the stride for full-volume prediction, are not stated 1612.05362 §2 1612.05362 §4.2.

Papers

• Medical Image Synthesis with Context-Aware Generative Adversarial Networks: introduces a two-stage auto-context GAN for MRI-to-CT synthesis, with a 3D generator and a gradient difference loss.