Student-Teacher Feature Pyramid Matching for Anomaly Detection

In plain words

Anomaly detection means finding the images or pixels that look unlike normal examples, and in factories or inspection it is usually impossible to collect examples of every defect. This paper takes a different route. It starts with a network that already knows how to recognise natural images, the teacher. It then trains a second copy of the same network, the student, to produce the same internal features as the teacher, but only on normal, defect-free pictures.

The student learns to copy the teacher wherever the picture looks normal. On a defect it has never seen, the two networks disagree, and the size of that disagreement marks where the defect is. The method compares several layers at once, from coarse texture-level layers to more abstract ones, so that both small and large defects show up. Because the comparison is one forward pass per network, the result is a pixel-level map and an image-level score in one go.

TL;DR

• One-class anomaly detection on normal images only. A pretrained ResNet-18 teacher is distilled into a student with the same architecture §3.1.
• Training minimises the normalised feature distance at three pyramid levels (conv2_x, conv3_x, conv4_x) with equal weights §3.2 eq. 1–eq. 3.
• Test-time anomaly map is the product of upsampled per-level distance maps; the image score is its maximum §3.3 eq. 4.
• Author-reported MVTec AD results: 0.955 image AUC-ROC, 0.970 pixel AUC-ROC, 0.921 PRO (ResNet-18, mean over 15 categories) tab. 1 tab. 2.
• Official code exists and uses several settings the paper does not state (SGD momentum and weight decay, ImageNet normalisation). See Paper vs code.

Problem

Anomalies are rare and unpredictable, so supervised labelling of all defect types is seen as unrealistic §1. Earlier pixel-level methods based on reconstruction or density estimation degrade when reconstructions are poor or likelihoods are miscalibrated §1. Patch-based pretrained-feature methods (Napoletano et al.) and the fully-convolutional approach (Sabokrou et al.) are cited as either slow or coarse §2.2. The student-teacher framework of Bergmann et al. is cited as efficient but as losing information in a lightweight teacher and in single-scale transfer; scaling needs several student-teacher pairs §1. The paper assumes all regions of the training images are anomaly-free §3.2.

Method

Framework §3.1. A fixed teacher, pretrained on image classification (ResNet-18 on ImageNet in the experiments), feeds a student with the identical architecture. Both networks output feature pyramids. Only the first three bottom layer groups (conv2_x, conv3_x, conv4_x of ResNet-18) are used §4.2 fig. 2.

Training §3.2. For each level l and each spatial position (i, j), the teacher and student feature vectors are divided by their L2 norms. The loss at that position is half the squared L2 distance between the two normalised vectors eq. 1. The image loss averages over positions eq. 2. The total loss is a weighted sum over levels with weights alpha_l >= 0; all weights are set to 1 in all experiments eq. 3 §3.2. The student is updated by minimising the minibatch mean of the image losses. The teacher is never updated §3.2.

Test §3.3. Each level gives an anomaly map of size w_l x h_l built from the per-position loss eq. 1. Each map is bilinearly upsampled to the image size. The final map is the element-wise product of the upsampled maps eq. 4. The image anomaly score is the maximum of this map §3.3.

# Training loss for one minibatch  [§3.2] [eq. 1]–[eq. 3]
# Input:  images (B, 3, 256, 256); frozen teacher; student
# Output: scalar loss

def training_loss(images, teacher, student):
    # Feature pyramids; the teacher gets no gradient  [§3.2]
    teacher_maps = teacher(images)      # L maps, each (B, d_l, h_l, w_l)
    student_maps = student(images)      # same shapes as teacher_maps
    total = 0
    for level in range(L):
        # Unit-normalise the channel vector at each position  [eq. 1]
        teacher_unit = normalise(teacher_maps[level], dim=1)
        student_unit = normalise(student_maps[level], dim=1)
        # Half squared L2 distance per position  [eq. 1]
        distance = 0.5 * sum_sq(teacher_unit - student_unit, dim=1)
        # Mean over positions, then weighted sum over levels  [eq. 2] [eq. 3]
        total = total + alpha[level] * mean_over_positions(distance)
    # Mean over the minibatch  [§3.2]
    return mean_over_batch(total)

# Anomaly map and image score for one test image  [§3.3] [eq. 4]
# Input:  image (1, 3, 256, 256)
# Output: anomaly map (256, 256); image score (scalar)

def anomaly_map(image, teacher, student):
    # Forward the same image through both networks  [§3.3]
    teacher_maps = teacher(image)
    student_maps = student(image)
    result = 1
    for level in range(L):
        # Per-position distance of [eq. 1], at the level's resolution  [§3.3]
        level_map = half_sq_distance(
            teacher_maps[level], student_maps[level])
        # Bilinear upsampling to the input size  [§3.3]
        level_map = upsample_bilinear(level_map, size=(256, 256))
        # Element-wise product over levels  [eq. 4]
        result = result * level_map
    # Image score is the maximum of the map  [§3.3]
    image_score = max(result)
    return result, image_score

The paper gives no algorithm box; the two blocks above are reconstructed from §3.2–§3.3 and eq. 1–eq. 4. Figure numbering in this card follows the order of figures in the LaTeX source, since no compiled PDF was checked.

Hyperparameters

Results

Author-reported. Setup: MVTec AD, 15 categories, ResNet-18 teacher pretrained on ImageNet unless marked. Pixel-level metrics are PRO (top row per category) and pixel AUC-ROC (bottom row) tab. 1. Image-level metric is image AUC-ROC tab. 2. Columns are abbreviated: the omitted baselines SSIM-AE, AnoGAN, CNN-Dict, Cut-Paste and Patch-SVDD are in tab. 1. Dash = not reported.

Pixel-level tab. 1 (STAD = Bergmann et al.; PaDiM-R18 and SPADE as cited in the table; * marks extra pretrained data in the paper)

Image-level tab. 2 (average image AUC-ROC over 15 categories; * as above)

Wide-ResNet50x2 backbone §4.3. With the SPADE backbone the text reports 0.973 AUC-ROC and 0.923 PRO. The text does not say which level (image or pixel) the 0.973 refers to; pixel level is likely given the surrounding sentences.

Ablation: pyramid levels tab. 3 (mean over 15 categories; block 2 = conv2_x)

Ablation: teacher pretraining data tab. 4 (mean over 15 categories)

Ablation: training-set size tab. 5 (mean over 15 categories; Ours vs SPADE)

The paper also reports that the gap to SPADE on pixel-level results reaches up to 11.7% §4.3. The metric and categories behind that figure are not stated.

Paper vs code

Code: official repository gdwang08/STFPM at commit 2598a5e35fd02f2f9dcfd0f3e8249adc22320e59, GPL-3.0. Owner matches the first author's name, and the README calls it the official implementation. Searched for links in the paper source (none), then GitHub repository search for the title and "STFPM"; this is the only repository that fits the owner rule.

• Optimiser. Paper: SGD with learning rate 0.4 §4.2. Code: SGD with momentum 0.9 and weight decay 1e-4 main.py.
• Loss scale. Paper: half squared distance, so 1 - cos per position eq. 1. Code: full squared distance, with no 1/2 main.py. The code loss is therefore twice the paper's loss, which doubles the gradient at the same learning rate. Evaluation maps are scaled the same way main.py; the constant does not change AUC or PRO rankings.
• Checkpoint selection. Paper: lowest validation error with the Eq. 1 error §4.2. Code: mean of the test-time product anomaly map over the validation set, saving the lowest main.py.
• Map resolution. Paper: upsample each level to w x h, then take the product §3.3 eq. 4. Code: product at 64 x 64, then resize to 256 for scoring main.py main.py.
• Epochs. Paper: 100 §4.2. Code default: 100 main.py; README example: 200 README.md.
• Normalisation. Paper: not stated. Code: ImageNet mean and std main.py.
• Teacher and student modes. Teacher is in eval mode and under no_grad main.py main.py. Student is in train mode during training main.py and eval mode for testing main.py.
• Validation split. 80/20 matches §4.2 main.py.
• Reported numbers. Paper Table 1 mean PRO 0.921 tab. 1; README mean PRO 0.9283467. Paper mean pixel AUC-ROC 0.970 and image AUC-ROC 0.955; README 0.970174 and 0.955488 README.md. Paper carpet PRO 0.958 and AUC-ROC 0.988 tab. 1; README carpet pixel AUC 0.990292 and PRO 0.966061 README.md. The README says the released checkpoints reproduce its numbers. Those checkpoints were not downloaded here.

Pitfalls

• Stated ranges are wrong. The paper says ell in (0, 1) §3.2 and Omega in [0, 1] §3.3. For unit vectors, half the squared distance equals 1 - cos, so it lies in [0, 2]. A product of L such maps can exceed 1. Do not clip the map or use a fixed [0, 1] threshold. Code maps are unbounded.
• Do not implement the paper's loss literally without checking the scale. The code's loss is twice the paper's main.py. Choose one convention and state it. The paper does not say which scale its lr 0.4 refers to.
• Optimiser details are missing from the paper. Momentum, weight decay and the lack of a schedule come from code only main.py.
• Validation uses normal images only. Checkpoints are chosen by error on held-out normal images. No anomalous data is used for selection.
• BatchNorm and small final batches. Batch statistics are used in student training. The loader keeps the last partial batch main.py. The paper does not mention BatchNorm. Neither the teacher's BatchNorm statistics nor the student's BatchNorm mode are specified in the paper.
• Wasted compute in the code. ResNet18_MS3 runs every module, including layer4, although only layers 4 to 6 are kept main.py. Truncate after layer3 if you need speed.
• Mask rasterisation. Ground-truth masks are resized with OpenCV's default interpolation and then cast to bool, so any non-zero interpolated pixel becomes anomalous and masks grow slightly main.py.
• Deprecated numpy and torchvision APIs. np.bool is used [code: main.py#L64] evaluate.py; it is removed in NumPy 1.24 and later. models.resnet18(pretrained=...) is deprecated in newer torchvision.
• Hardware. The code calls .cuda() unconditionally main.py main.py.
• PRO is expensive. The compiled text defers the PRO definition to [Bergmann2020]. The code sweeps 4000 thresholds with Python loops over connected components evaluate.py evaluate.py; the README warns it "might take a long time" README.md.
• Few-shot subsets. Table 5 uses 5% and 10% of training images §4.3; the sampling procedure is not stated.
• Teacher data. Only ImageNet is used for the main results; the other datasets appear only in the ablation tab. 4. The paper's own text says that generalisation from MNIST and SVHN is poor §5.3.
• Test cost. Each test image needs two ResNet-18 forward passes and an upsampling at full resolution for each level.
• Figure and table numbers are taken from source order because no compiled PDF was checked.

Lineage

• Builds on: the student-teacher anomaly detector of Bergmann et al. (CVPR 2020, "Uninformed Students"; cited as STAD). This paper argues that STAD uses two-step, single-scale transfer and loses information §4.3. It is also a case of feature-based knowledge distillation [Wang & Yoon 2020, arXiv:2004.05937] §3.1. No card exists for either.
• Compared against: SPADE (Cohen & Hoshen, arXiv:2005.02357) §4.3 tab. 1 tab. 2; GANomaly, ITAE, Cut-Paste, Patch-SVDD and PaDiM tab. 1 tab. 2; STAD tab. 1.
• Superseded by: not known from content already in this wiki.