PaperlayerView exact revision on arXiv

Source-linked AI summary

Differentiable Patch Selection for Image Recognition

Jean-Baptiste Cordonnier, Aravindh Mahendran, Alexey Dosovitskiy, Dirk Weissenborn, Jakob Uszkoreit, Thomas Unterthiner

arXiv:2104.03059v1cs.CVcs.AIcs.LGstat.ML

TL;DR

Processing high-resolution images in full is costly, while downsampling can discard fine details. This paper selects task-relevant patches with an end-to-end differentiable Top-K method, and reports results across traffic-sign recognition, inter-patch reasoning, and fine-grained recognition. In the billiard-ball reasoning task, its transformer and max-pooling variants solve the task with high accuracy and outperform the other methods.

  • Problem

    Full-resolution image processing is computationally expensive, while downsampling may remove important fine details and hurt performance.

    Source: §1 · p. 1

  • Method

    A scorer ranks patches for differentiable Top-K selection, after which feature and aggregation networks process the selected patches end-to-end.

    Source: §1 · p. 2, §3 · p. 3

  • Results

    On the billiard-ball task, transformer and max-pooling aggregation solve the task with high accuracy and outperform the other methods.

    Source: §5.2 · p. 7

  • Takeaways & Limitations

    The method supports arbitrary downstream neural processing of selected patches without auxiliary losses for training the scorer.

    Source: §6 · p. 8

    Patch selection can struggle to account for global context, and scorer training can become unstable or lock onto wrong patches.

    Source: §6 · p. 8

Abstract

from arXiv · show

Neural Networks require large amounts of memory and compute to process high resolution images, even when only a small part of the image is actually informative for the task at hand. We propose a method based on a differentiable Top-K operator to select the most relevant parts of the input to efficiently process high resolution images. Our method may be interfaced with any downstream neural network, is able to aggregate information from different patches in a flexible way, and allows the whole model to be trained end-to-end using backpropagation. We show results for traffic sign recognition, inter-patch relationship reasoning, and fine-grained recognition without using object/part bounding box annotations during training.

1. Introduction

High-resolution images are costly to process in full, while downsampling can remove details needed for recognition. The paper frames patch selection as a task-dependent ranking problem and proposes an end-to-end trainable approach for tasks where relevant regions are sparse.

1. Introduction

  • Full-resolution processing is computationally expensive, but downsampling can remove fine details and hurt performance.The goal is to save compute while retaining the ability to recognize fine details.Source: §1 · p. 1
  • Because relevant regions are task-dependent and patch-retain decisions are discrete, selecting them directly is difficult to train end-to-end.For traffic signs, for example, a system may ignore sky but must not miss signs.Source: §1 · p. 1
  • The method ranks patches with a small ConvNet and selects the Top K for downstream processing, training selection through backpropagation.It is most effective when most patches are irrelevant and their locations are not known in advance.Source: §1 · p. 1
  • The approach targets image-recognition settings rather than image coverage, such as semantic segmentation or object detection.The paper does not aim to preserve coverage of the full image.Source: §1 · p. 2
  • The paper evaluates the method on street-sign recognition, synthetic inter-patch reasoning, and fine-grained classification without bounding-box annotations.The stated settings omit object/part box annotations during both training and evaluation for fine-grained classification.Source: §1 · p. 2

2. Related Work

Prior approaches either do not reduce the detail at which all image regions are processed, rely on difficult sequential attention, or constrain how patch information is combined. The paper’s differentiable Top-K selection is presented as a way to sample patches while allowing more flexible aggregation.

2. Related Work

  • Region-proposal methods select regions for later processing, but the cited object-detection architectures require bounding-box supervision, unlike this paper’s weakly supervised image-classification setting.The paper focuses on training from a single class label per image.Source: §2 · p. 2
  • Soft-attention masks typically still process the whole image at fixed resolution, so they do not provide computational-efficiency gains.Other cited approaches use multiple resolutions or auxiliary losses to extract patches.Source: §2 · p. 2
  • Attention-based multiple-instance learning can identify relevant patches, but processes all patches in equal detail and therefore does not use selection to save computation.The passage describes this as common in medical imaging with very large images.Source: §2 · p. 2
  • Sequential glimpse methods can be difficult to train with reinforcement learning, while spatial-transformer gradients can be very local; this method computes gradients with respect to all patches at each backward step.The cited alternative by Angles et al. uses an auxiliary function to lift non-differentiable Top-K.Source: §2 · p. 2
  • Unlike ATS’s average-pooling restriction, differentiable Top-K permits flexible combination of information from selected patch embeddings.ATS samples patches and averages their embeddings.Source: §2 · p. 2
  • The method uses perturbed optimizers to differentiate Top-K and reports better performance than the Sinkhorn operator in its experiments.The cited alternatives include Gumbel-Softmax and optimal-transport ranking or sorting formulations.Source: §2 · p. 2

3. Model overview

The model scores image regions, selects a subset, computes a feature embedding for each selected patch, and aggregates those embeddings into a prediction. Its components allow different downstream processing and aggregation choices.

3. Model overview

  • The model’s pipeline comprises a scorer, patch selector, feature network, and aggregation network, with learnable scorer, feature, and aggregation parameters.The scorer’s selected patches are processed by a downstream network, and the model is trained end-to-end without an ad hoc selection loss.Source: §3 · p. 2, Figure 2 · p. 3
  • The scorer assigns task-relevance scores to image regions, and the selection module extracts the K highest-scoring patches for feature processing.The scorer is intended to be shallow and may operate on downscaled images; the model typically selects few regions, such as 10, rather than the 2000 used in Faster-RCNN.Source: §3 · p. 3
  • A feature network embeds each selected patch, and an aggregation network combines the embeddings into the prediction.Aggregation may be mean or max pooling, a transformer, or another operation; the architecture permits different choices.Source: §3 · p. 3
  • The framework distinguishes patch scoring and extraction from feature computation and aggregation, unlike architectures that merge feature extraction and aggregation into a transformer over all patches.Vision transformers described here exhaustively extract all P × P patches with stride P.Source: §3 · p. 3

4. Patch selection as differentiable Top-K

The paper formulates patch selection as sorted Top-K and makes it differentiable with the perturbed maximum method. Noise-perturbed selections are averaged during training to provide gradients across patches, while inference uses hard Top-K.

4. Patch selection as differentiable Top-K

  • Patch selection is formalized as Top-K so the full model can train end-to-end without auxiliary losses for individual components.The Top-K indices identify the K largest scores.Source: §4 · p. 3

4.1. Patch selection as Top-K

  • Sorted Top-K represents chosen indices as one-hot indicators, enabling selected patches to be extracted from the tensor of all patches; this discrete operation is non-differentiable.Sorting is imposed because arbitrary output permutations would break the perturbed-optimizer method.Source: §4.1 · p. 4

4.2. Differentiable Top-K

  • The perturbed maximum method makes patch selection differentiable by perturbing the input, solving the linear program for each perturbation, and averaging the results.The experiments use an empirical mean over 500 independent samples and tune σ over {0.01, 0.05, 0.5}.Source: §4.2 · p. 4
  • The linear-program constraints give each assignment column total weight one and impose sorted indices, with an index-sorted integer Top-K solution among its optima.The paper notes that the linear program has infinitely many optimal solutions.Source: §4.2 · p. 4
  • Averaging perturbed selections can yield weighted combinations of patches early in training and gives gradients that account for all image patches.Sorting keeps patch assignments aligned across perturbed inputs; the authors suggest image mixing may contribute to generalization.Source: §4.2 · p. 4
  • Hard Top-K is used at inference for efficiency, while σ is decayed to zero during training to bridge the train-test gap.At σ = 0 the differentiable operation matches hard Top-K numerically, but scorer gradients vanish.Source: §4.2 · pp. 4–5

5. Experiments and results

The experiments assess patch selection on traffic-sign recognition, inter-patch reasoning, and fine-grained bird classification. Results are competitive on traffic signs, strongest with flexible aggregation on the billiard-ball task, and slightly below NTS-Net on CUB-200.

5.1. Traffic Signs Recognition

  • On traffic signs, our method matched ATS’s reported mean performance but had larger variance; ATS’s 0.01 entropy setting had 20–25% of runs below approximately 60% test accuracy.The setup used 747 training and 684 test images, with 2,028 candidate patches and K∈{5, 10}.Source: Table 1 · p. 5, §5.1 · p. 5

5.2. Inter-patch reasoning

  • On the billiard-ball task, Transformer and max-pooling aggregation achieved high accuracy and outperformed other methods; a standard CNN took 24% longer to train than Top-K with a Transformer.Mean-pooling and Transformer runs were relatively stable, while some ATS and max-pooling runs predicted the majority class.Source: §5.2 · pp. 6–7, Table 2 · p. 7

5.3. Fine-grained bird classification

  • On CUB-200, the best mean-aggregation model reached 87.3%, yet remained slightly below NTS-Net; overlapping patches around the bird head may worsen overfitting.Source: Table 3 · p. 7, §5.3 · p. 8
  • For CUB-200, the method omits NTS-Net’s ranking loss and uses Squeeze-Excitation scorer layers to learn patch selection without hard greedy NMS.Mean aggregation averages per-patch embeddings and adds them to the whole-image embedding, yielding 2048-dimensional aggregated embeddings.Source: §5.3 · p. 8

6. Discussion and future work

The paper introduces differentiable Top-K patch selection for end-to-end, task-dependent learning and reports competitive results on localized classification and improved state-of-the-art performance on inter-patch reasoning. It also identifies training instability, difficulty accounting for global context, and dataset-dependent overfitting as areas for improvement.

6. Discussion and future work

  • The differentiable Top-K module enables end-to-end, task-dependent patch selection and supports arbitrary downstream processing without auxiliary scorer losses.The authors report competitive performance on localized classification and a significant state-of-the-art advance on inter-patch relationship reasoning.Source: §6 · p. 8
  • Patch selection has a chicken-and-egg difficulty: choosing optimal patches requires image contents, while efficiently learning those contents requires choosing the right patches.The authors connect this difficulty to accounting for global context.Source: §6 · p. 8
  • Training can be unstable, with majority-class predictions or confident selection of wrong patches emerging early or during training.The authors suggest more adaptive optimization might improve learning stability.Source: §6 · p. 8
  • Overfitting differs by dataset: patch selection reduced it on Traffic street signs, but the scorer memorized CUB-200 watermarks and backgrounds; NTS-Net reduced it considerably.The authors suggest larger datasets or self-supervised learning as possible future directions.Source: §6 · p. 8

B. Speed Improvements by Sampling Patches

On a single V100, processing selected patches can speed inference while retaining a ResNet50 feature network. The reported gain depends on the useful image content being localized, and hard Top-K is consistently faster than differentiable Top-K at inference.

B. Speed Improvements by Sampling Patches

C. Linear Program

The linear program represents Top-K selection with one-hot assignments and an objective that sums selected scores. Its strict inequality constraint also enforces ordering by index.

C. Linear Program

  • Each column of an integral LP solution is a one-hot indicator selecting one score from the patch scores.Source: C. Linear Program · p. 11
  • Because the objective sums selected scores, the optimal integral solution selects the K highest scores.Source: C. Linear Program · p. 11
  • The strict inequality constraint forces lower-index selections to come earlier; otherwise swapping their order contradicts validity.Source: C. Linear Program · p. 11

D. Decaying σ to 0

On the billiard balls task with transformer aggregation, hard Top-K at inference and linearly decaying perturbation magnitude each improve accuracy in the reported ablation. Figure 9 compares constant perturbation settings with a setting that decays σ to zero.

D. Decaying σ to 0

E. Differentiable Sinkhorn

The Sinkhorn Top-K approach follows Xie et al. [44] and achieves good traffic-sign results, but performance suffers when discrete Top-K is used at inference.

E. Differentiable Sinkhorn

  • The Sinkhorn Top-K implementation follows Algorithm 3 of Xie et al. [44] and gives good traffic-sign results, but suffers with discrete Top-K at inference.The reported results use a setting similar to Section 5.1; Table 5 reports mean and standard deviation across five runs.Source: E. Differentiable Sinkhorn · p. 12, Table 5 · p. 12

F. Experimentation Details

The experiments use a four-layer convolutional scorer and dataset-specific feature networks; optimization, augmentation, and the CUB-200 adaptation vary by dataset.

F. Experimentation Details

  • The scorer uses four 3 × 3 convolutions with stride 1 and valid padding, with 8, 16, 32, and 1 feature maps, followed by 8 × 8 max pooling.The first three convolutions use ReLU activations.Source: F. Experimentation Details · p. 12
  • Feature extractors are a narrow ResNet18 for Swedish traffic signs, standard ResNet18 for billiard balls, and ResNet50 for CUB-200.The traffic-sign network matches ATS [22], while the CUB-200 extractor matches NTS-Net.Source: F. Experimentation Details · p. 12
  • Traffic-sign and billiard-ball experiments use ADAM-W, while CUB-200 uses SGD with momentum 0.9; weight decay with momentum reduced CUB-200 overfitting.The NTS-Net data augmentation is retained, with pixels rescaled to [−1, 1] to match ImageNet-pretrained ResNet50 inputs; squeeze-excitation layers are added to the CUB-200 region proposal network to compensate for lacking non-maximum suppression.Source: F. Experimentation Details · p. 12

G. Qualitative Results

On CUB-200, the selected patches often cover the bird’s eyes and torso, though they can overlap; the appendix also shows how the differentiable Top-K operator is implemented.

G. Qualitative Results

  • The K = 4 CUB-200 visualization achieves 87.3% top-1 accuracy and often selects bird-eye and torso regions, with substantial patch overlap.The visualization is from the best-performing of five random seeds, and selected anchors extending beyond image boundaries cause padding artifacts.Source: G. Qualitative Results · p. 12, Figure 11 · p. 14
  • The adapted NTS-Net region proposal head inserts a squeeze-excitation layer before anchor prediction to communicate across features, then applies zero-one normalization across 1614 scores.Figure 10 contrasts the baseline and adapted architectures.Source: Figure 10 · p. 13

H. Differential Top-K PyTorch Code

  • The PerturbedTopK operator accepts k, num_samples, and sigma, then adds sampled Gaussian noise to scores before selecting the top k values.The code forms perturbed score vectors and calls torch.topk on each; for Gaussian noise, the noise and gradient are the same.Source: H. Differential Top-K PyTorch Code · p. 15
  • Its backward pass combines perturbation-based gradients with the output gradient to compute the input gradient.The code uses einsum operations for both gradient aggregation steps.Source: H. Differential Top-K PyTorch Code · p. 15

Related papers

Loading 2104.03059v1…