[Paper Review] DISK: Learning local features with policy gradient

Michał J. Tyszkiewicz, Pascal Fua+1 more|arXiv (Cornell University)|Jun 24, 2020

Advanced Image and Video Retrieval Techniques50 references184 citations

TL;DR

DISK learns local keypoints and descriptors end-to-end using reinforcement learning, optimizing for many correct feature matches with a discrete-keypoint policy. The approach yields dense, discriminative features and achieves state-of-the-art results on three benchmarks.

ABSTRACT

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.

Motivation & Objective

• Motivate the challenge of learning local features end-to-end due to discreteness in keypoint selection and matching.
• Propose a reinforcement learning-based framework (DISK) to optimize for correct feature matches.
• Create a simple, expressive probabilistic model that aligns training and inference regimes for reliable from-scratch training.
• Demonstrate that dense, discriminative features can outperform traditional assumptions about keypoint quality.

Proposed method

• Propose a probabilistic model that handles discrete keypoint selection for RL optimization.
• Use policy gradient optimization to maximize a surrogate objective of many correct feature matches.
• Train end-to-end to couple keypoint detection, orientation, and description while maintaining differentiability.
• Leverage RL signals to avoid relying on inexact objectives and non-differentiable components.
• Enable dense feature extraction with strong discriminative power and robust downstream matching.

Experimental results

Research questions

• RQ1 Can a discrete-keypoint policy be optimized end-to-end via reinforcement learning to maximize correct matches?
• RQ2 Does a dense, discriminative feature representation learned through policy gradient achieve superior matching performance?
• RQ3 How does DISK compare to prior hand-crafted or partially learned descriptors in end-to-end training settings?
• RQ4 Can training from scratch with a surrogate RL objective yield reliable convergence for local feature learning?
• RQ5 Do the learned features maintain performance across multiple public benchmarks?

Key findings

• DISK achieves state-of-the-art results on three public benchmarks.
• The features can be extracted very densely while remaining discriminative.
• The method overcomes training challenges associated with non-differentiable keypoint selection.
• Training from scratch with a surrogate objective yields robust learning and reliable convergence.
• The approach relaxes some limitations of prior RL-based local feature methods by not relying on non-differentiable end-to-end supervision.