Source-linked AI summary
DISK: Learning local features with policy gradient
Michał J. Tyszkiewicz, Pascal Fua, Eduard Trulls
TL;DR
Sparse local-feature systems are difficult to train end to end because keypoint selection and matching are discrete and computationally complex. DISK uses a probabilistic policy-gradient framework that samples keypoints and optimizes rewards for correct matches, with analytical gradients enabling training from scratch. It achieves state-of-the-art results across stereo and 3D reconstruction benchmarks, including first place in the Image Matching Challenge’s 2k-keypoint category.
Problem
End-to-end local-feature learning is limited by discrete keypoint selection and matching, whose joint computation is costly and whose feature quality depends on the selected feature set.
Method
DISK uses CNN-generated heatmaps and descriptors, samples discrete keypoints, probabilistically relaxes cycle-consistent matching, and applies policy gradients to maximize geometric-match rewards.
Results
DISK achieves state-of-the-art pose accuracy for stereo and 3D reconstruction and places #1 in the 2k-keypoint Image Matching Challenge category.
Takeaways & Limitations
DISK can train from scratch while producing many more matches than competing methods, supporting dense yet discriminative local features for downstream reconstruction tasks.
Takeaways & Limitations
The grid-based method allows at most one feature per cell and ignores neighboring cells, so non-maxima suppression at inference introduces a potentially sub-optimal training–inference mismatch.
Abstract
from arXiv · showhide
Local feature frameworks are difficult to learn in an end-to-end fashion, due to the discreteness inherent to the selection and matching of sparse keypoints. We introduce DISK (DIScrete Keypoints), a novel method that overcomes these obstacles by leveraging principles from Reinforcement Learning (RL), optimizing end-to-end for a high number of correct feature matches. Our simple yet expressive probabilistic model lets us keep the training and inference regimes close, while maintaining good enough convergence properties to reliably train from scratch. Our features can be extracted very densely while remaining discriminative, challenging commonly held assumptions about what constitutes a good keypoint, as showcased in Fig. 1, and deliver state-of-the-art results on three public benchmarks.
1 Introduction
DISK addresses the difficulty of end-to-end local-feature learning by using policy gradients to optimize discrete keypoint selection and matching for correct geometric correspondences. Its probabilistic modeling improves convergence while narrowing the gap between training and inference.
- Motivation: End-to-end local-feature learning is difficult because joint detection and matching are discrete, computationally complex, and dependent on the other selected features.Matching two feature sets costs O(|FA|·|FB|), while a feature’s usefulness depends partly on whether other features are similar to it.
- Method: DISK uses a CNN to produce keypoint heatmaps and dense descriptors, samples discrete keypoints, and builds a distribution over matches.Geometric ground truth supplies rewards for sampled matches, enabling optimization of expected per-match reward by gradient descent.
- Method: Careful probabilistic modeling yields analytical gradient expressions instead of noisy Monte Carlo approximations, allowing models to train reliably from scratch.This design narrows the gap between training and inference while retaining policy-gradient expressiveness.
- Contribution: In the SfM example, DISK produces associations on seemingly textureless regions where SIFT fails, yielding more landmarks with more observations.The figure compares Upright Root-SIFT and DISK using a reconstruction from 1179 images.
- Contribution: The approach is end-to-end trainable, targets correct feature matches directly, and improves downstream performance compared with earlier methods.The paper illustrates this benefit through Structure-from-Motion results and reports more accurate matches.
2 Related Work
Prior learned local-feature methods commonly target individual stages or use end-to-end relaxations with approximations, broken differentiability, or strong assumptions. Related RL approaches address detection or nondifferentiable objectives but face domain, supervision, or initialization constraints.
- Existing local-feature methods: Traditional local-feature pipelines separate keypoint detection, orientation estimation, and description, while learned methods often improve only one of these components.Descriptor learning on predefined keypoints remains competitive alongside methods focused on keypoints or orientations.
- End-to-end methods: Other end-to-end pipelines optimize the full process but may use inexact objectives, break differentiability, or assume descriptor-space extrema are good features.These design choices limit how directly the training objective represents the intended feature-selection process.
- Reinforcement-learning approaches: GLAMpoints uses RL to find keypoints correctly matched by SIFT descriptors, relying on deterministic matching and Q-learning for domain-specific retinal-image requirements.Its approach addresses detection rather than jointly learning general descriptors and matching.
- Reinforcement-learning approaches: Reinforced Feature Points learns with a general nondifferentiable objective involving RANSAC, but a single reward creates weak training signals and requires pretrained models for fine-tuning.DISK instead relaxes this setup by training a surrogate objective focused on finding many correct feature matches.
3 Method
DISK models feature extraction and matching probabilistically, then trains the pipeline with policy gradients using rewards for correct matches. Its design preserves a close relationship between training and inference while enabling exact gradient computation for the matching step.
- Probabilistic formulation: DISK represents feature sets and image matches as conditional probability distributions, enabling reinforcement-learning optimization of the feature pipeline.Feature extraction samples from detection heatmaps, while matching is modeled conditionally on the extracted feature sets.
- Feature extraction: A U-Net produces a detection map K and N descriptor channels, with N=128 chosen for comparison with SIFT and modern descriptors.The resulting feature set pairs sampled locations with l2-normalized descriptors at those locations.
- Feature extraction: Within each h × h grid cell, DISK samples a location using softmax preference and accepts it using a sigmoid-based absolute-quality probability.The total sampling probability combines relative preference among cell locations with the absolute detection quality of the proposed pixel.
- Inference: At inference, DISK replaces probabilistic sampling with arg max and a sign-based acceptance rule, then uses cycle consistency, a validation-selected ratio test, and NMS.NMS addresses grid-cell blindness and the one-feature-per-cell restriction, but introduces a potential training–inference misalignment.
- Matching: Matching relaxes cycle consistency by sampling forward and reverse descriptor matches and accepting pairs only when both samples are consistent.The forward and reverse categorical distributions are derived from the descriptor-distance matrix and a temperature parameter θM.
- Policy-gradient training: Because match probabilities are exact and rewards factorize over individual matches, DISK computes exact matching gradients without sampling, reducing gradient-estimation variance.The reward assigns positive or negative values to correct and incorrect matches, while non-matching features receive zero reward and a small keypoint penalty regularizes excess features.
4 Experiments
DISK is evaluated on Image Matching Challenge, HPatches, and ETH-COLMAP, with experiments covering implementation choices, downstream performance, and training–inference differences.
- Training data: The experiments use 135 MegaDepth scenes containing 63k posed images, with COLMAP poses and dense depth for pixel correspondences.Test-overlapping scenes are omitted, and co-visible image pairs are sampled using a heuristic.
- Implementation: The feature extractor is a 1.1M-parameter U-Net variant with a 219 × 219-pixel formal receptive field.It uses four down- and up-blocks, instance normalization, and PReLU nonlinearities.
- Optimization: Training processes three co-visible images per batch and accumulates rewards and gradients over all three image pairs.The matching parameter is gradually increased with a fixed schedule while the feature network is optimized by gradient descent.
- Image Matching Challenge: On the 2k Image Matching Challenge category, DISK outperforms all methods by 9.4% relative in stereo and 6.7% relative in multiview.On 8k, it outperforms all baselines on average across stereo and multiview but ranks slightly below the top three submissions.
- HPatches: On HPatches, DISK achieves state-of-the-art performance, ranking first on viewpoint scenes and second on illumination scenes.Combined, it outperforms its closest competitor, RFP, by 12% relative, despite training on non-planar data without strong affine transformations.
- Ablations and discussion: Smaller training grids improve performance but increase extracted features, distance-matrix size, and computational expense.Inference-time NMS is clearly beneficial despite departing from the training pipeline, while 19.9% of grid-selected features have a neighbor within 2 px.
5 Conclusions and future work
DISK is a probabilistic, end-to-end local-feature learner trained with policy gradient that can train from scratch and achieves state-of-the-art pose accuracy for stereo and 3D reconstruction.
- DISK learns local features end to end with a probabilistic policy-gradient approach.
- The method can train from scratch and yields many more matches than competing methods.
- DISK achieves state-of-the-art pose accuracy for stereo and 3D reconstruction.
- DISK ranks first in the 2k-keypoints category of the Image Matching Challenge using off-the-shelf matchers.
- Future work will replace the match relaxation with learned matchers.
Broader impact
The authors do not expect new, specific issues from their method, but identify potential misuse and broader inequities associated with keypoint-based automation and data access.
- The authors do not expect their method to create new, specific issues beyond those associated with existing keypoint applications.
- Potential concerns include visually guided missiles and localizing photographs without user consent, which could further compromise web privacy.
- Automation of data processing may disproportionately benefit established players with greater access to data and resources.
APPENDIX — DISK: Learning local features with policy gradient
The appendix identifies the supplementary material accompanying the NeurIPS submission.
- The appendix is supplementary material for NeurIPS submission 1194.
Training data
Training pairs are selected by COLMAP-based 3D-keypoint co-visibility, filtered to a specified ratio range, and organized into triplets across 135 scenes.
- Training image pairs use COLMAP 3D keypoints to compute a co-visibility proxy and exclude pairs that are not co-visible or too easy.
- Pairs are retained when the co-visibility ratio satisfies 0.15 ≤ r ≤ 0.8.
- Triplets are formed by sampling a seed image and two associated images, without requiring the latter pair to meet the co-visibility criterion.
- Scenes overlapping Image Matching Challenge test or validation subsets are manually blacklisted.
- Scenes with low-quality depth maps are also blacklisted, and scenes producing fewer than 10k co-visible triplets are removed.
- The final training set contains 135 scenes and approximately 133k co-visible triplets.
Continuous evaluation
The authors pause training every 5k optimization steps, corresponding to 10k triplets, to evaluate stereo performance continuously.
- Every 5k optimization steps, the authors evaluate stereo performance on a smaller validation subset.The interval corresponds to 10k triplets.
- They re-implement the benchmark's mAA(10o) metric for this evaluation.
- The best model is selected using this metric before hyper-parameter tuning proceeds.
Computational cost
Inference runs at approximately 7 frames per second for 1024 × 1024 inputs, while training takes approximately 1.2 seconds per two 768 × 768 triplets.
- ≈7 frames per second are obtained at inference for 1024 × 1024 inputs.
- Training with 768 × 768 inputs requires ≈1.2 seconds per two triplets.
- Gradient accumulation can reduce memory requirements and enable training on 12 Gb GPUs.
Qualitative results for epipolar supervision – Fig. 7
The paper compares depth-based and epipolar supervision qualitatively: depth supervision usually avoids textureless areas, whereas epipolar supervision places points near 3D object boundaries and on thin structures.
- Qualitative results for epipolar supervision – Fig. 7: Epipolar-based supervision produces points around the boundaries of 3D objects.
- Qualitative results for epipolar supervision – Fig. 7: Depth-based supervision usually teaches models to avoid textureless areas such as the sky.
- Qualitative results for epipolar supervision – Fig. 7: Thin structures such as lamp-posts create features with epipolar supervision but not with depth supervision.
- Qualitative results for epipolar supervision – Fig. 7: Some epipolar-supervised points may remain unmatched, as illustrated by the obelisk in the rightmost images for (c) and (d).
- Qualitative results for epipolar supervision – Fig. 7: The scene breakdown reports results for 9 scenes and their average, using 2k and 8k feature settings.