Few-Shot Classification with Feature Map Reconstruction Networks

In plain words

A few-shot classifier has to recognise new categories from only a handful of labelled pictures. Common methods compare one summary vector per picture. The summary is usually an average over the whole picture, which throws away the small details that tell two similar birds or aircraft apart. Keeping every spatial position, on the other hand, makes the comparison sensitive to where an object sits in the frame.

This paper takes a different route. For a new picture, it tries to rebuild each of its spatial positions out of the features of the example pictures from one class, using any position in those examples. If the class is right, the rebuild is good; if the class is wrong, it is poor. The class with the best rebuild wins. The rebuild is a standard least-squares fit that has a closed-form answer, so no iterative solver runs at test time, and it adds only a few learned numbers to the network.

The authors report that this beats several strong baselines on fine-grained bird, aircraft and species datasets, and is competitive on the two standard ImageNet-derived benchmarks. These are their own measurements.

TL;DR

FRN (Feature Map Reconstruction Networks) scores each class by the negative mean squared error of a ridge-regression reconstruction of the query feature map from the pooled support feature maps of that class §3.1 eq. 1–eq. 5. The regularisation and a recalibration factor are learned, with three learned parameters in total §3.2. The authors report the best numbers on four fine-grained benchmarks, and mixed results against baselines on mini-ImageNet and tiered-ImageNet tab. 3 tab. 4 tab. 5 tab. 6 tab. 7. The released code differs from the paper in the auxiliary loss, the pre-training penalty and two small constants; see Paper vs code.

Problem

Global average pooling discards spatial detail that fine-grained tasks need §1. Flattening the feature map keeps the detail but encodes object location, which is nuisance information §1. Larger receptive fields overfit to background cues §1. Earlier reconstruction methods either solve an iterative optimal-transport problem (DeepEMD) or use learned attention modules with many extra parameters (CrossTransformer, CrossAttention) §2. The paper's aim is a reconstruction-based classifier that keeps spatial detail, ignores location, and needs no extra modules or test-time gradient steps §1 §2 §3.7.

Method

Episode setup. An n-way, k-shot episode has n classes, k labelled support images per class and q query images per class §2. The backbone turns each image into a feature map Q with r spatial positions and d channels. Conv-4 uses d = 64 and r = 5×5 = 25; ResNet-12 uses d = 640 and r = 25 at 84×84 input §4 [SM Network Backbones].

Support pool. For each class c, the feature vectors of all k support images are stacked across images and locations into one matrix S_c with kr rows and d columns §3.1.

Ridge reconstruction. For a query feature map Q, FRN solves the ridge problem that finds the r-by-kr weight matrix W which best rebuilds Q from S_c, with a penalty λ on W §3.1 eq. 1. The solution is closed form eq. 2 [eq. 3]. The paper gives the form (S^T S + λI)^(-1) S^T S for the reconstruction (Woodbury form) and uses it in all experiments §3.4 eq. 9 [SM Pseudo-Codes].

Class score. The score for class c is the negative mean over locations of the squared error between Q and its reconstruction §3.1 eq. 4. Class probabilities are a softmax over the negative scores, multiplied by a learned temperature γ §3.1 eq. 5.

Learned regularisation. The penalty is scaled by kr/d, so that a support pool concatenated with itself gives the same reconstruction §3.2. The learned form is λ = (kr/d)·e^α and the reconstruction is recalibrated by ρ = e^β, with α and β initialised to zero §3.2 eq. 6. The full prediction is ρ·Q·(S^T S + λI)^(-1)·S^T S eq. 7 eq. 8. Together with γ, this is the only set of learned parameters added beyond the backbone §3.2.

Loss. The query probabilities of each episode go through cross-entropy, as in standard episodic meta-training §3.2.

Auxiliary loss. A term L_aux, scaled by 0.03, encourages support features of different classes to be orthogonal, after row normalisation §3.5 eq. 12. The authors say it has little or no effect on FRN §3.5 [SM Auxiliary Loss and Regularization Parameters].

Pre-training (ResNet-12, mini-ImageNet and tiered-ImageNet). A learnable r×d matrix M_c is kept for every base class c and acts as that class's support pool, so the pool holds one map §3.6. Reconstruction error stays the logit, α and β are fixed at zero, and γ stays learnable §3.6 eq. 10 eq. 11. The M_c matrices are discarded after pre-training, so the model size is unchanged §3.6. Pre-training is followed by episodic fine-tuning §4.2 §5.1.

Inference. There is no gradient step at test time. Each query is assigned to the class with the smallest reconstruction error §4.2 [SM Computational Efficiency].

# Negative reconstruction error for every query and class  [§3.1] [eq. 9]
# Input:  queries (m, r, d); supports (n, k*r, d), pooled per class
# Output: scores (m, n), larger is more likely

def reconstruction_scores(queries, supports, alpha, beta):
    # Penalty lambda = (k*r / d) * e^alpha  [§3.2] [eq. 6]
    penalty = (supports.shape[1] / supports.shape[2]) * exp(alpha)
    # Recalibration rho = e^beta  [§3.2] [eq. 6]
    rho = exp(beta)
    # Gram matrix of each class pool, (n, d, d)  [eq. 9]
    gram = transpose(supports) @ supports
    # Closed-form solve (S^T S + lambda I)^-1 S^T S  [eq. 9]
    solved = inverse(gram + penalty * identity(d)) @ gram
    # Reconstructed query maps, (n, m*r, d)  [eq. 9]
    recon = rho * (queries @ solved)
    # Squared error summed over channels, (n, m, r)  [§3.1] [eq. 4]
    errors = square(recon - queries).sum(axis=-1)
    # Mean over locations, then negate  [§3.1] [eq. 4]
    return -errors.mean(axis=-1).transpose()

# One meta-training episode  [§3.1] [§3.2] [§3.5]
# Input:  images of n classes (k support, q query each); query labels
# Output: scalar loss

def episode_loss(images, query_labels, alpha, beta, gamma):
    # Feature maps from the backbone, (n*(k+q), r, d)  [§3.1] [§4]
    maps = backbone(images)
    # Pool support maps per class (n, k*r, d) and keep query maps  [§3.1]
    supports, queries = split_episode(maps)
    # Scores scaled by the learned temperature gamma  [§3.1] [eq. 5]
    logits = gamma * reconstruction_scores(queries, supports, alpha, beta)
    # Cross-entropy over the query predictions  [§3.2]
    frn_loss = cross_entropy(logits, query_labels)
    # Auxiliary term with weight 0.03  [§3.5] [eq. 12]
    aux_loss = 0.03 * auxiliary_term(supports)
    return frn_loss + aux_loss

# Non-episodic pre-training head  [§3.6] [eq. 10] [eq. 11]
# Input:  images (B, 3, 84, 84); labels over all base classes
# Output: scalar loss

def pretrain_loss(images, labels, gamma, dummy_maps):
    # Feature maps of the batch, (B, r, d)  [§3.6]
    maps = backbone(images)
    # One dummy (r, d) map per base class, used as its support pool  [§3.6]
    # alpha = beta = 0 are fixed; gamma stays learnable  [§3.6]
    scores = reconstruction_scores(maps, dummy_maps, 0.0, 0.0)
    return cross_entropy(gamma * scores, labels)

Hyperparameters

Benchmark-specific values come from the supplementary training section. The paper says one hyperparameter set is used for all FRN models except where noted [SM Hyperparameters]. Values marked "code" are not stated in the paper.

Results

Author-reported. Setup: 5-way episodes, 10,000 test episodes, 95% confidence intervals, unless the row says otherwise. FRN results are from models trained with 5-shot episodes §5.1. Baseline rows are the authors' re-implementations where marked † in the table, and published numbers otherwise §4.1.

Ablation, mini-ImageNet, ResNet-12:

Analysis points from the paper: 1-shot-trained FRN is worse than 5-shot-trained FRN even at 1-shot evaluation §5.1 [SM Episodic Training Shot Number]. The auxiliary loss and the learned λ and ρ change results only a little, and not consistently across backbones [SM Auxiliary Loss and Regularization Parameters]. Training memory for CTX rises from 6225 MB (5×5, 1-shot) to 9693 MB (10×10, 5-shot), while FRN at 10×10 uses 6727–7125 MB depending on the form and shot [SM Memory comparisons]. Hardware for the latency and memory measurements is not stated.

Paper vs code

Code: Tsingularity/FRN at 64e43f3ca3a4953445788fc853e288e032d0083c (MIT, copyright Luming Tang). The paper links the repository in its abstract.

• Auxiliary loss, normalisation. Paper: features are row-normalised onto the unit sphere §3.5. Code divides the whole support tensor by one Frobenius norm over all classes and locations frn_train.py.
• Auxiliary loss, pair sum. Paper: the double sum covers all ordered pairs i, j with j ≠ i §3.5 eq. 12. Code loops over unordered pairs only (j < i), which halves the sum relative to the paper's formula frn_train.py.
• Auxiliary loss, weight. Both use 0.03 §3.5 frn_train.py.
• Penalty constant. Code adds 1e-6 to λ, which the paper does not show FRN.py.
• Pre-training penalty. The paper's pre-training equation writes λI with no stated pool size §3.6 eq. 10. Code treats each dummy map as a pool of one map, so the penalty is r/d (with α = 0) FRN.py [code: models/FRN.py#L137].
• Temperature γ. Paper: learnable, initial value not stated §3.1. Code: initialised to 1.0 FRN.py.
• Meta-iNat learning-rate schedule. The supplement says the rate is halved "every 20 episodes" [SM Fine-Grained Benchmarks]. The script halves it every 20 epochs, for 100 epochs in total train.sh.
• Validation trials. Not stated in the paper. Code uses 1000 validation episodes trainer.py.
• Checked and consistent. Reconstruction is the Woodbury form (eq. 9) FRN.py. The distance averages over locations, not a sum FRN.py. Test queries are 16 per class and the evaluation uses 10,000 trials eval.py [code: trainers/trainer.py#L245]. Cut points for CUB, mini-ImageNet and tiered settings match the supplement train.sh train.sh. The Conv-4 backbone and the ResNet-12 leaky ReLU (0.1) match ResNet.py.

Pitfalls

• Auxiliary loss is ambiguous as written. The paper's normalisation and pair sum are not what the code does (see Paper vs code). The paper reports little effect of dropping it §3.5, so the difference is likely small, but an implementation should pick one form and say which.
• Train on 5-shot episodes. Models trained on 1-shot episodes are worse on 1-shot evaluation too §5.1. The main results use 5-shot-trained models §4 [SM Episodic Training Shot Number].
• Pre-training is needed for mini- and tiered-ImageNet, not for the fine-grained benchmarks. The paper says pre-training is not used in fine-grained experiments for fairness §3.6 §4.1. An FRN trained from scratch beats a pre-trained FRN evaluated without episodic fine-tuning §5.1.
• Average-pool pre-training does not help FRN. The paper says the resulting features are not spatially distinct enough, and the classifier is worse than training from scratch §5.1 tab. 9. FRN needs its own pre-training head §3.6.
• Scale the ResNet-12 embedding by √640. The paper says this is necessary for stable training [SM Hyperparameters]. Its image decoder, by contrast, runs without this normalisation, because the regularised reconstructions otherwise fall off the decoder's input manifold [SM Image Decoder].
• Cost scales with d, not only with kr. The Woodbury form inverts a d×d matrix, once per class per episode §3.4. For ResNet-12 (d = 640), the paper's own memory table shows the other form using less memory at 5-shot 5×5, against its stated analysis [SM Computational Efficiency].
• Score is a mean over locations. The squared error is averaged over the r locations §3.1 FRN.py. The paper's kr/d penalty scaling is what keeps reconstructions unchanged when a support pool is repeated §3.2.
• Do not mix tiered-ImageNet releases. The 84 px release (Ren et al.) and DeepEMD's 224 px release give different numbers, and both are in one table §4.2 tab. 7. Use the same release when comparing.
• Inference preprocessing differs by dataset. CUB raw and mini-ImageNet resize the shorter edge to 92 then centre-crop 84; cropped CUB and Aircraft resize to 92×92 [SM Pre-processing] util.py. The paper says these choices are inherited and were not tuned.
• Cross-domain numbers depend on the CUB class split. The paper reports three splits separately §4.3 tab. 8. Do not average across them.
• Some margins are small. On mini-ImageNet, FRN's 1-shot result is below FEAT (66.45 vs 66.78) tab. 7. On 84 px tiered-ImageNet, FRN is below RFS-distill on both columns tab. 7. The paper's wording "highly competitive" §4.2 covers these.
• Numerical precision of the matrix inverse is not stated. The paper gives no precision or regularisation floor beyond λ. The code adds 1e-6 to λ and sets no dtype for the inverse FRN.py [code: models/FRN.py#L77].
• Pre-training uses all base classes. The dummy maps cover every base class (for example 64 on mini-ImageNet), not the episode classes §3.6.

Lineage

• Builds on: Prototypical Networks (Snell et al., 2017) for episodic training §3.2; the learned temperature of Gidaris et al. (2018), Chen et al. (2020) and Ye et al. (2020) §3.1; and the Woodbury form used by Bertinetto et al. (2018) §3.4. Pre-training follows Chen et al. (2020) and Ye et al. (2020) §3.6.
• Compared against: DeepEMD (Zhang et al., 2020), an optimal-transport reconstruction; CrossTransformer (Doersch et al., 2020) and CrossAttention (Hou et al., 2019), attention-based reconstruction; Deep Subspace Networks (Simon et al., 2020), closed-form projection. The paper re-implements CTX and DSN as baselines §3.7 [SM Implemented Baselines].
• Superseded by: not stated in this card.