Segment Anything Model is a Good Teacher for Local Feature Learning

Our SAMFeat uses a backbone for feature extraction, along with several heads serving different purposes end-to-end. The simple overview of SAMFeat is shown in Figure 2 and the detailed network structure is shown in Figure 3.

We utilize a weighted binary cross-entropy loss for the detector due to the significant imbalance between keypoints and non-keypoints. With the predicted keypoint heatmap $H\in\mathbb{R}^{h\times w}$ and the pseudo-ground truth label $G\in\mathbb{R}^{h\times w}$, the detector loss $\mathcal{L}_{\text{det}}$ is defined as follows:

where the weight $\lambda$ is empirically set to 200.

The attention map proves to be highly useful in descriptor optimization, distillation optimization, and matching processes [45]. In terms of descriptor and distillation optimization, we provide a detailed analyze of motivation in Section 2 and Section 3. For matching, attention-weighted local descriptors are more effective. Regions with high attention scores in one image are likely to match with similar high-score regions in another image, reducing the matching space and enhancing accuracy. These consistent attention scores serve as prior information, making local descriptor matching more efficient and reliable. The Attention Head will be jointly optimized during the optimization process of both descriptor generation and feature distillation.

where $\|\cdot\|$ represents the L2 norm and $(i,j)$ denotes the index position of the local descriptor $d$.

Inspired by [45, 29], the descriptor loss splits the optimization into two parts: the descriptor’s angle and the consistent attention weight. This differs from the standard triplet loss, which only focuses on the angle component. Positive samples have converging attention scores, while negative samples diverge, leading to a consistent distribution of attention scores across image pairs. Higher attention scores significantly influence the gradient, allowing the network to prioritize more relevant samples and avoid optimizing descriptors for less informative pixels, like the sky or grass.

The descriptor Loss is formulated to jointly optimize local descriptors and consistent attention. Building on previous research [29, 45], we use the corresponding point sets $(P,P^{\prime})$ obtained from ground-truth camera parameters and depths to supervise the training of local descriptors. For an image pair $(I,I^{\prime})$, dense descriptors $D,D^{\prime}$ and attention maps $A,A^{\prime}$ are extracted using our SAMFeat method.

Given a point set $P$ of size $M$ in image $I$ and the corresponding points $P^{\prime}$ in image $I^{\prime}$, the local descriptors of $P,P^{\prime}$ are denoted by Equation 3 as $d_{i}$ respectively, where $i\in\{1,\ldots,M\}$. The corresponding score of the descriptor $d_{i}$ on the attention map $A$ is denoted as $\nu_{i}$. Thus, the attention-weighted descriptor is defined as $y_{i}=\nu_{i}\cdot d_{i}$. For $y_{i}$, its positive distance $\|y_{i}\|^{+}$ is defined as:

and its hardest negative distance $\|y_{i}\|^{-}$ is defined as:

The overall descriptor loss is the sum of the individual losses:

where $\nu$ is the attention score corresponding to $y$, and $T$ is a smoothing factor that adjusts the effect of attention weighting on the loss. $T$ is set to $15$ following [45].

SAM [11] is a newly released visual foundation model for segmenting any objects and has strong zero-shoot generalization due to the fact that it is trained using 11 million images and 1.1 billion masks. Due to its scale, model distillation [46] is deployed in this work. We freeze the weights of SAM and use its output as pseudo-ground truth to guide more accurate and robust local feature learning. In this subsection, we introduce how we utilize three gifts from SAM to enhance our SAMFeat. Shown in Figure 3, we input the image $I$ into the SAM [11] with frozen parameters and then simply processed to produce the following three outputs for guided local feature learning.

Pixel-wise Representations Relationship: SAM’s image encoder trained from 11 million images is used to extract image representations for assigning semantic labels. The representation of the encoder outputs implies a valuable semantic correspondence, i.e., pixels of the same semantic object are closer together. To eliminate the effect of specific semantic categories on generalizability, we adopt relations between representations as distillation targets. SAM’s encoder outputs $\mathcal{F}\in\mathbb{R}^{\frac{1}{8}H\frac{1}{8}W\times C}$, where $C$ is the channel number for feature map. The pixel-wise representations relationship can be defined as $\mathcal{R}\in\mathbb{R}^{\frac{1}{8}H\frac{1}{8}W\times\frac{1}{8}H\frac{1}{8}W}$, where $\mathcal{R}(i,j)=\frac{\mathcal{F}(i)\cdot\mathcal{F}(j)}{|\mathcal{F}(i)||\mathcal{F}(j)|}$.

Semantic Grouping: We use the automatically generating masks function¹¹1https://github.com/facebookresearch/segment-anything/ of SAM to obtain fine-grained semantic groupings. Specifically, it works by sampling single-point input prompts in a grid over the image, and SAM can predict multiple masks from each of them. Then, masks are filtered for quality and deduplicated using non-maximal suppression [11]. The semantic grouping of the output can be defined as ${G}\in\mathbb{R}^{H\times W\times N}$, where $N$ is the number of semantic groupings. Notice that semantic grouping differs from semantic segmentation in that each grouping does not correspond to a specific semantic category (e.g. buildings, car, and person).

Edge Map: The binary edge map ${E}\in\mathbb{R}^{H\times W\times 1}$ is derived directly ²²2https://github.com/ymgw55/segment-anything-edge-detection from the segmentation results of SAM, which highlights the fine-grained object boundaries.

Thanks to the gifts of the foundation model, SAM, we are able to consider SAM as a knowledgeable teacher with intermediate products and outputs to guide the learning of local features. First, we employ Attention-weighted Semantic Relation Distillation (ASRD) to distill the category-agnostic semantic relations in the SAM encoder into SAMFeat, thereby enhancing the expressive power of local features by introducing semantic distinctiveness. Second, we utilize the high-level semantic grouping of SAM outputs to construct Weakly Supervised Contrastive Learning Based on Semantic Grouping (WCS), which provides cheap and valuable supervision for local descriptor learning. Third, we design an Edge Attention Guidance (EAG) to utilize the low-level edge structure discovered by SAM to guide the network to pay more attention to these edge regions, which are more likely to be detected as keypoints and rich in discriminative information during local feature detection and description.

However, not every pixel is equally important for local feature learning, and forcing the network to learn a large proportion of background pixels (e.g., the sky) can hinder the learning of discriminative foreground regions. We therefore advocate a greater focus on discriminative foreground regions in the distillation process. In contrast to classical relational distillation [47], we propose Attention-weighted Semantic Relation Distillation (ASRD) to guide the distillation process to focus on valuable relational pairs to further motivate the transfer of knowledge that facilitates local feature matching.

Specifically, $C^{o}_{4}$ is exported from Conv7 layer and then imported into the distillation head to get $C^{d}_{4}\in\mathbb{R}^{\frac{1}{8}H\times\frac{1}{8}W\times 128}$. Following the operations reported in Sec. III-B, the semantic relation matrix of $C^{d}_{4}$ can be defined as $\mathcal{R}^{{}^{\prime}}$. As shown in Figure 4, we distill the semantic relation matrix by imposing L1 loss in order to obtain semantic discriminativeness for $C^{d}_{4}$. $\mathcal{R}^{{}^{\prime}}$ and $\mathcal{R}$ are the corresponding student (SAMFeat) and teacher (SAM) relation matrix. As shown in Fig. 4 (b), we use the attention map obtained in Section 2 to construct the weight matrix used to weight the relational distillations. We flatten the attention map into a 1-dimensional vector $V_{A}$, and then multiply $V_{A}$ and the transpose $V_{A}^{T}$ of $V_{A}$ by matrix multiplication to obtain the weight matrix $\mathcal{W}$, which can be formally defined as follows:

where the elements of $\mathcal{W}\in\mathbb{R}^{\frac{1}{8}H\times\frac{1}{8}W}$ and $\mathcal{R}$ correspond to each other. Attention-weighted Semantic Relation Distillation Loss $\mathcal{L}_{dis}$ can be defined as:

where $N$ is the number of matrix elements, i.e., $(\frac{1}{8}H\times\frac{1}{8}W)\times(\frac{1}{8}H\times\frac{1}{8}W)$. Since ASRD is category-agnostic, it is possible to generalize local feature distillation semantic information to generic scenarios. A detailed pseudo-code is illustrated in algorithm 1.

Here ${\rm dis(P_{i},P_{j})}$ means calculate the Euclidean distance between the local descriptors corresponding to the two sampling points $P_{i}$ and $P_{j}$. $G(\cdot)$ denotes the indexed semantic grouping category. $J$ denotes the number of positive samples, noting that since $J$ is not consistent for each image, we take the average to denote the positive sample distance. Similarly, the negative sample average distance $D_{neg}$ can be defined as:

where ${\rm M}$ is a margin parameter used to protect distinctiveness within semantic groupings, and ${\rm T}$ means the temperature coefficient.

Figure 8 visualize the learning outcome of object boundaries under the guidance of SAM. With our EAG, SAMFeat learns accurate boundaries and edges effectively and efficiently. This prior knowledge of boundaries and edges will then aid feature detection and description.

With accurate loss supervision, SAMFeat is able to utilize the learned knowledge in later modules discussed in section 3 to further aid feature learning and matching tasks.

To generate our training data with dense pixel-wise correspondences, we rely on the MegaDepth dataset [10], a rich resource containing image pairs with known pose and depth information from 196 diverse scenes. Specifically, we use MTLDesc [29] ³³3https://github.com/vignywang/MTLDesc released megedepth image and the correspondence ground truth for training. In our experiment, we meticulously configured the parameters to establish a consistent and efficient training process. Hyper-parameters are set as follows. The learning rate of 0.001 enables gradual parameter updates, and the weight decay of 0.0001 helps control model complexity and mitigate overfitting. With a batch size of 14, our model processes 14 samples per iteration, striking a balance between computational efficiency and convergence. $\rm M$ and $\rm T$ are set to 0.07 and 5. Training spans 30 epochs to ensure comprehensive exposure to the data, with a total training time of 3.5 hours. By meticulously defining these parameters and configurations, we establish a robust experimental setup that ensures replicability and accurate evaluation of our model’s performance. More detailed information about parameter tuning and ablation experiments can be found in the supplementary material.

We evaluate the performance of our method in the image-matching tasks on the most popular feature learning and matching benchmark: HPatches [1]. The HPatches dataset consists of 116 sequences of image patches extracted from a diverse range of scenes and objects. Each image patch is associated with ground truth annotations, including key point locations, descriptors, and corresponding homographies. We follow the same evaluation protocol as in D2-Net [48], where eight unreliable scenes are excluded. To ensure an equitable comparison, we align the features extracted by each method through nearest-neighbor matching. A match is deemed accurate if its estimated reprojection error is lower than a predetermined matching threshold. The threshold is systematically varied from 1 to 10 pixels, and the mean matching accuracy (MMA) across all pairs is recorded, indicating the proportion of correct matches relative to potential matches. Subsequently, the area under the curve (AUC) is computed at 5px based on the MMA. The comparison between SAMFeat and other state-of-the-art methods on HPatches image matching is visualized in Figure 7. The MMA @3 threshold against other state-of-the-art methods under each threshold is listed in Table I. SAMFeat achieved the highest MMA @3 even compared to the most updated feature learning model in 2023 top-tier conferences.

To further validate the efficacy of our approach when dealing with intricate tasks, we assess its performance in the area of visual localization. This task involves estimating the camera’s position within a scene using an image sequence and serves as an evaluation benchmark for local feature performance in long-term localization scenarios, without requiring a dedicated localization pipeline. We utilize the Aachen Day-Night v1.1 dataset [60] to showcase the impact on visual localization tasks. To ensure fairness in the assessment, we employ a predefined visual localization pipeline⁴⁴4https://github.com/GrumpyZhou/image-matching-toolbox based on colmap provided by benchmark⁵⁵5https://www.visuallocalization.net. This pipeline operates as follows: Initially, custom features extracted from the database’s images are employed to construct a structure-from-motion model. Subsequently, the query images are registered within this model using the same custom features. For keypoints matching, we utilize the mutual nearest neighbor approach to effectively filter out outliers. And we set the number of features in this experiment to $10000$. We tally the number of accurately localized images under three distinct error thresholds, namely (0.25m, 2°), (0.5m, 5°), and (5m, 10°), signifying the maximum allowable position error in both meters and degrees. Referring to Table II, we categorize current state-of-the-art methods into two categories: $G$ contains methods that are designed for general feature learning tasks; $L$ contains methods that are designed, tuned, and tested specifically for localization tasks, and they typically perform poorly outside of specific localization scenarios, as shown in Table I. SAMFeat achieved the top performance among all general methods, while also revealing a competitive performance among methods that are designed specifically for visualization.

We utilize the ETH benchmark [63] to evaluate the performance of our method on the 3D reconstruction task. Three medium-scale datasets from the benchmark are used for this purpose. We perform exhaustive image matching on all three collections and apply ratio test filtering with a default threshold of 0.8 to eliminate incorrect matches. The reconstruction protocol follows the COLMAP [2] pipeline, which involves running Structure-from-Motion (SfM) first and then Multi-View Stereo (MVS) to generate the dense point cloud models. As shown in Table III, SAMFeat consistently produces the highest number of registered images and sparse points, along with competitive track length and reprojection error, demonstrating its robustness and accuracy in 3D reconstruction tasks.

We conduct ablation studies on different aspects to support our claim and illustrate the necessity of each of our contributed modules.

Even though each additional loss function introduces some extra training time, the tradeoff is justified by the minimal increase in time and the corresponding improvement in accuracy. The final training time remains acceptable, demonstrating the lightweight nature of our approach. This efficiency makes our method easily implementable and resource-efficient, highlighting its reproducibility and practicality for individual researchers in the field of feature learning and description.

To further demonstrate SAMFeat’s high efficiency and fast inference speed, we conduct a comparison of inference time between other state-of-the-art methods in table X. We assessed the running speed of various methods using open-source code. In Table X, our approach demonstrated exceptionally competitive performance while maintaining a fast inference speed among many lightweight methods.

As mentioned in Section III-B in the full paper, SAMFeat learns edge maps from SAM and utilizes Edge Attention Guidance (EAG) to further enhance the precision of local feature detection and description by encouraging the network to prioritize attention to the edge region. Figure 8 demonstrates the learning outcome of SAMFeat. With the fine-grained object boundaries from SAM, SAMFeat is able to learn clear object edges. This illustrates two things: first, the encoded feature that is used to generate the edge map contains rich edge information, and second, with a clear and accurate generated edge map and EAG, SAMFeat is able to better capture the details of edge areas and improve the robustness of local descriptors.