Cost aggregation

Overview

A cost volume holds a score for every candidate match. In few-shot segmentation, it scores every position in a query image against every position in one or more support images 2207.10866 §3.3. In stereo matching, it scores each left-image pixel against right-image pixels shifted by each candidate disparity 1903.04025 §3.2. Raw scores of this kind are noisy: repeated textures and clutter produce false matches 2207.10866 §1. Cost aggregation is the step that mixes information across the volume, so that each score can use context from neighbouring positions before a decoder reads out a mask or a disparity 2207.10866 §1 1903.04025 §3.3.

Two families appear here. The few-shot aggregators treat the volume as a 4D table of query and support positions. The stereo aggregator treats it as a 3D volume of disparity, height and width 1903.04025 §3.3. The readouts differ: a mask decoder in the first case 2207.10866 §3.5, a soft argmin over disparity in the second 1903.04025 §3.4, eq. 4. The group-wise volume that the stereo card builds is a construction step, not an aggregator, so it is only the input here 1903.04025 §3.2.

The variants differ in how they split the volume into tokens or kernels, how far a token or kernel can reach, and what that costs in memory and time.

Variants

Tokenisation of the volume

A transformer needs tokens. Patch-style embedding cuts the volume into non-overlapping 4D windows. The paper argues that this cuts context at window edges and gives the model weak convolutional bias 2207.10866 §3.4.1 §3.4.2. The alternative stacks overlapping 4D convolutions, each followed by pooling, normalisation and a ReLU, and projects to a 128-channel embedding 2207.10866 §3.4.2 §4.1. In the paper's FSS-1000 ablation, the convolutional front end raises 1-shot mIoU for each aggregator: linear transformer 87.0 to 87.7, Fastformer 87.1 to 87.8, Swin 89.9 to 90.3 2207.10866 app. B.

Standard transformer over the flattened volume

Not usable in the paper's setting. The volume runs out of memory for the mIoU and mBA rows, and the memory column reports about 84 GB for the standard transformer 2207.10866 tab. 5. The batch size is not stated.

Centre-pivot 4D convolution

The aggregator used by HSNet, per the paper's table. It reports 88.1 mIoU, 3.5 GB and 52.7 ms per run on one RTX 3090 2207.10866 tab. 5. It has the lowest run-time in that table, 4.6 ms below the proposed method.

Linear transformer

Reports 87.7 mIoU, 3.5 GB and 56.8 ms 2207.10866 tab. 5.

Fastformer

Reports 87.8 mIoU, 3.5 GB and 122.9 ms 2207.10866 tab. 5. The same table lists it as the slowest of the runnable aggregators.

4D convolutional Swin transformer

The paper's own choice. The volume is split into 4D windows of size 4 along each axis, self-attention runs inside each window, and the next block shifts the windows by half a window 2207.10866 §3.4.1 §4.1. The aggregator runs coarse to fine over a three-level pyramid, and each level is added back to its input as a residual 2207.10866 §3.4.2 §3.4.3. Reports 90.3 mIoU, 68.0 mBA, 3.8 GB and 57.3 ms 2207.10866 tab. 5. Card: Cost Aggregation with 4D Convolutional Swin Transformer for Few-Shot Segmentation.

The paper's component ablation on FSS-1000 adds the pieces one at a time. The baseline with a 2D convolutional decoder is at 80.0 mIoU (1-shot). Adding the Swin transformer gives 85.4, and adding the convolutional front end gives 87.0 2207.10866 tab. 4. The backbone for that table is not stated.

Stacked 3D hourglass over a disparity volume

The aggregator of the group-wise stereo network. A pre-hourglass module of four 3D convolutions, each with batch normalisation and ReLU, adds its input back as a residual. Three stacked 3D hourglass networks follow [arxiv:1903.04025 §3.1] §3.3 tab. 1. Each hourglass downsamples twice, to 2C and then 4C channels, and upsamples back with transposed convolutions. Its shortcuts use 1×1×1 convolutions, which the paper says add little computation 1903.04025 §3.3 tab. 1. After the pre-hourglass module and after each hourglass, a one-channel output module gives a score per disparity. A softmax over disparity and a soft argmin then give the disparity [arxiv:1903.04025 §3.4]. At inference the paper drops the first three output modules 1903.04025 §3.3.

The stereo ablations isolate two choices. Replacing the original hourglass with the new one lowers Scene Flow EPE from 0.876 to 0.808 px, and run-time falls from 241.0 to 198.3 ms on 480×640 inputs 1903.04025 tab. 2. The group count changes the volume this aggregator receives. With the base aggregator, 40 groups give EPE 1.212 px, against 1.369 px for one group, which is full correlation 1903.04025 tab. 2. Card: Group-wise Correlation Stereo Network.

When to use which

Timings are not comparable across the two papers. The few-shot numbers are from one RTX 3090 and an i7-10700, for the paper's setup 2207.10866 §4.4, and the shot setting for 2207.10866 tab. 5 is not stated. The stereo numbers are for 480×640 inputs on one TITAN Xp 1903.04025 tab. 2. The stereo paper does not report memory use. Neither paper's timings carry over to other hardware.

Implementation sketch

Language-agnostic outline of the shared core, following the paper's pyramid. The residual on the attention output follows the paper's text and released code; its equation for the encoder omits it 2207.10866 §3.4.2.

# Coarse-to-fine cost aggregation over a correlation pyramid  [§3.4.3]
# corr_levels: list of 4D volumes, coarsest first

def aggregate(corr_levels):
    prev = None
    for corr in corr_levels:
        # Overlapping 4D convolutions as the tokeniser  [§3.4.2]
        x = volumetric_conv(corr)
        if prev is not None:
            # Add the coarser result, upsampled, as guidance  [§3.4.3]
            x = x + upsample_bilinear(prev)
        # Windowed 4D attention, added back as a residual  [§3.4.1] [§3.4.2]
        prev = x + window_attention_4d(x, window=4)
    return prev

The stereo variant keeps the same shape but convolves over disparity and produces one score per disparity at each output stage. Channel count 64 is the concatenated group-wise and concatenation volume 1903.04025 tab. 1.

# Stacked 3D hourglass aggregation over a stereo cost volume  [§3.3] [§3.4]
# volume: (C, D, H, W), with C = 64  [tab. 1]

def aggregate_stereo(volume):
    # Pre-hourglass: four 3D convolutions with a residual add  [§3.1] [tab. 1]
    cost0 = pre_hourglass(volume)
    # Three stacked hourglasses, each fed by the previous one  [§3.3]
    cost1 = hourglass(cost0)
    cost2 = hourglass(cost1)
    cost3 = hourglass(cost2)
    # One score per disparity, softmax over D, then soft argmin  [§3.4, eq. 4]
    stages = [cost0, cost1, cost2, cost3]
    return [soft_argmin(softmax(score_head(c), over=D)) for c in stages]

Contradictions

• Augmentation. The paper says no augmentation is used [arxiv:2207.10866 §4.1]. Its appendix reports PASCAL-5i mean mIoU of 67.9 without augmentation, 67.3 with one photometric set, 66.1 with another and 65.7 with both. FSS-1000 stays within 90.0 to 90.3 2207.10866 app. B. The released configs enable augmentation, and the card's Paper vs code section records this as unresolved.
• K-shot voting. The paper divides the vote sum by K before thresholding at 0.5 2207.10866 §3.6. The released code divides by the per-image maximum vote count instead, so the threshold is relative to that count. Both are recorded in the card's Paper vs code section; which one produced the reported numbers is not stated.
• Checkpoint selection. The released training script keeps the checkpoint with the best mIoU on the evaluation fold. The paper does not state this, so the per-fold numbers in 2207.10866 tab. 1 may be optimistic. See the card's Paper vs code section.
• Hourglass gain. The text says the new hourglass "increases EPE by 7.8%" on Scene Flow 1903.04025 §4.4. Table 2 shows EPE falling from 0.876 to 0.808 1903.04025 tab. 2. The stereo card treats the table as the reliable source; the text is not resolved here.
• Test crop. The paper tests on full 960×540 Scene Flow images 1903.04025 §4.2. The released code crops to 960×512, so the reported numbers may come from the smaller input. See the stereo card's Paper vs code section.

No other disagreements are known in this wiki.

Papers

• Cost Aggregation with 4D Convolutional Swin Transformer for Few-Shot Segmentation: replaces the fixed 4D convolutions over the correlation volume with a windowed 4D transformer, with a coarse-to-fine pyramid and an appearance-aware decoder.
• Group-wise Correlation Stereo Network: builds the stereo cost volume from channel groups and aggregates it with a 3D stacked hourglass network that outputs a disparity per stage.