A nearest neighbor retrieval method based on automatic construction of neighbor graph index

Abstract

The present invention discloses a nearest neighbor retrieval method based on the automatic construction of a nearest neighbor graph index, which extracts and fuses vector data sets and graph convolutional neural networks respectively through a self-attention mechanism feature extraction module and a graph convolutional neural network. k The characteristics of the neighbor graph use the parameter performance prediction module to predict the performance indicators according to the output of the feature extractor module and the splicing result of the parameter configuration, thereby completing the construction of the neighbor graph index. Compared with the existing neighbor graph parameter selection method, the model proposed in the present invention can improve the accuracy and efficiency of neighbor graph parameter selection, and can meet the vector retrieval requirements of related applications in large-scale scenarios, and optimize the user experience; at the same time, when extracting features from vector data sets, the present invention uses a graph neural network for feature extraction. Compared with traditional feature engineering methods, the use of a graph neural network can extract richer features, thereby making the model more accurate when predicting neighbor graph parameter performance.

Description

Technical Field

[0001] The present invention belongs to the technical field of vector retrieval, and in particular relates to a nearest neighbor retrieval method based on automatic construction of a nearest neighbor graph index. Background Art

[0002] With the development of artificial intelligence, vector retrieval technology has been widely used in applications such as face recognition, file search, and speech recognition. Among them, vector retrieval based on neighbor graph has become a research hotspot in academia and industry due to its excellent search and navigation capabilities.

[0003] However, when building a neighbor graph index, it is necessary to enumerate all possible parameter configurations through grid search, which is a very time-consuming operation. To solve this problem, one solution is to use machine learning methods to predict the performance of a set of neighbor graph parameters and then filter out the best performing parameters. However, since the potential features of the vector data set cannot be effectively extracted during the feature extraction process of the vector data set and the method itself is complex, the accuracy and efficiency of predicting the performance of the neighbor graph parameters will be reduced. Summary of the invention

[0004] The object of the present invention is to provide a nearest neighbor retrieval method based on automatic construction of a nearest neighbor graph index.

[0005] In a first aspect, the present invention provides a nearest neighbor retrieval method based on automatic construction of a nearest neighbor graph index, which comprises the following steps:

[0006] Step 1: obtain multiple sample data sets, perform feature extraction on the multiple sample data sets respectively, and obtain multiple vector data sets; the multiple vector data sets form a set V;

[0007] Step 2: construct a k-nearest neighbor graph for each vector data set in the set V; set multiple sets of neighbor graph parameter configurations for each k-nearest neighbor graph, and construct a neighbor graph index; obtain the performance index corresponding to each neighbor graph index; construct a training set, the training set includes multiple samples, each sample includes a vector data set, a k-nearest neighbor graph and its parameter configuration and performance index;

[0008] Step 3, establish a nearest neighbor graph index automatic construction model; the nearest neighbor graph index automatic construction model includes a feature extractor module and a parameter performance prediction module; the feature extractor module extracts and fuses the features of the vector data set and the k nearest neighbor graph as the output of the feature extractor module; the input of the parameter performance prediction module is the splicing result of the output of the feature extractor module and the parameter configuration; the output of the parameter performance prediction module is the predicted performance index;

[0009] Step 4: Use the training set constructed in step 2 to train the neighbor graph index automatic construction model;

[0010] Step 5: Automatically build a model based on the trained neighbor graph index to obtain the optimal prediction performance index of the sample data set corresponding to the query data set; select the parameter configuration corresponding to the optimal prediction performance index to generate the neighbor graph index; after vectorizing the query data set, perform neighbor retrieval on the generated neighbor graph index.

[0011] Preferably, the feature extractor module includes a self-attention mechanism feature extraction module and a graph convolutional neural network; the self-attention mechanism feature extraction module is used to extract features from the vector data set; and the graph convolutional neural network is used to extract features from the k-nearest neighbor graph.

[0012] Preferably, the parameter performance prediction module adopts a multi-layer perceptron.

[0013] Preferably, the sample data set is an image data set or a text data set; the images in the image data set are feature extracted using a convolutional neural network to obtain image vectors; the texts in the text data set are extracted using a pre-trained language model to obtain text vectors.

[0014] Preferably, the performance indicators include recall rate, number of queries processed per second and number of distance calculations.

[0015] Preferably, the neighbor graph parameter configuration includes the neighbor graph index type, the maximum number of neighbors, the number of nearest neighbors and the size of the candidate set when composing the graph.

[0016] Preferably, the method for setting multiple groups of parameter configurations is: enumerating different combinations of neighbor graph index type, maximum number of neighbors, number of nearest neighbors, and size of candidate set when composing a graph.

[0017] Preferably, in step 4, the neighbor graph index automatically constructs the loss function of the model The expression is:

[0018]

[0019] Among them, x i Automatically construct the input of the model for neighbor graph indexing; i is the actual performance; f(x i ) is the performance of the model prediction automatically constructed by the neighbor graph index; n is the number of groups of parameter configurations in the input model.

[0020] In a second aspect, the present invention provides a computer device comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the memory stores the computer program; and the processor executes the above-mentioned nearest neighbor retrieval method automatically constructed based on a neighbor graph index.

[0021] In a third aspect, the present invention provides a readable storage medium storing a computer program; when the computer program is executed by a processor, it is used to implement the above-mentioned nearest neighbor retrieval method automatically constructed based on a neighbor graph index.

[0022] The present invention has the following beneficial effects:

[0023] 1. The present invention predicts neighbor graph performance indicators by establishing a neighbor graph index automatic construction model, thereby selecting parameter configuration and completing the construction of the neighbor graph index. Compared with the existing neighbor graph parameter selection method, the model proposed in the present invention can improve the accuracy and efficiency of parameter configuration selection, and can meet the vector retrieval requirements of related applications in large-scale scenarios, optimizing the user experience.

[0024] 2. When extracting features from vector data sets, the present invention uses a graph neural network for feature extraction. Compared with traditional feature engineering methods, the use of a graph neural network can extract richer features, thus making the model more accurate when predicting the performance of neighbor graph parameters. BRIEF DESCRIPTION OF THE DRAWINGS

[0025] Figure 1 A flowchart for automatically constructing a model for training a neighbor graph index in the present invention.

[0026] Figure 2 A flowchart of automatically building a model using a neighbor graph index in the present invention. DETAILED DESCRIPTION

[0027] The present invention will be further described below in conjunction with the accompanying drawings.

[0028] A nearest neighbor retrieval method based on automatic construction of a nearest neighbor graph index comprises the following steps:

[0029] Step 1: Use an image dataset or a text dataset as a sample dataset; obtain multiple sample datasets, each of which has a label; perform feature extraction on multiple sample datasets to obtain a set of vector datasets. ; Where M is the number of vector data sets; is the i-th vector data set, ; For vector dataset The jth vector in; N is the vector data set The number of vectors in the is the dimension of each vector; For vector dataset Corresponding data distribution space. In this embodiment, the images in the image data set are feature extracted using a convolutional neural network (CNN) to obtain image vectors that express the visual content of the image; the texts in the text data set are processed using the BERT pre-trained language model to generate text vectors of words, sentences, and paragraphs containing context information.

[0030] Step 2: For each vector data set in set V Construct a k-nearest neighbor graph and set multiple neighbor graph parameter configurations; neighbor graph parameter configurations include neighbor graph index type , Maximum number of neighbors , the number of nearest neighbors and the size of the candidate set L when composing the graph; enumerate all parameter configurations through grid search, build the corresponding neighbor graph index, and obtain the performance indicators of the neighbor graph index; the performance indicators of the neighbor graph index include recall (recall rate), qps (number of queries processed per second) and distcomp (number of distance calculations); build a dataset for training the model, which includes a set V of vector datasets and each vector dataset Corresponding multiple k-nearest neighbor graphs, parameter configurations, and performance metrics.

[0031] In this embodiment, during the process of enumerating all parameter configurations, the neighbor graph index type ; Maximum number of neighbors ;Number of nearest neighbors ; The candidate set size during composition is L={40,50,60,70,80,90,100,110,120,130,140,150}.

[0032] Step 3: Figure 1 As shown in the figure, a neighbor graph index automatic construction model is established; the neighbor graph index automatic construction model includes a feature extractor module and a parameter performance prediction module; the feature extractor module includes a self-attention mechanism feature extraction module and a graph convolutional neural network (GCN); the self-attention mechanism feature extraction module adopts a graph transformer; the graph transformer considers the vector data set through a global attention mechanism All-pair attention between vectors in the graph convolutional neural network; the input of the graph convolutional neural network is the k-nearest neighbor graph; the output of the graph transformer and the graph convolutional neural network are fused to obtain a vector dataset The embedding of is the output of the feature extractor module;

[0033] The parameter performance prediction module adopts a multi-layer perceptron (MLP); the input of the parameter performance prediction module includes the embedding of the vector data set extracted by the feature extractor and the parameter configuration enumerated for the vector data set; the parameter prediction performance module predicts the performance indicators of each set of parameter configurations in the vector data set by combining the embedding and parameter configuration of the vector data set.

[0034] Step 4: Use the dataset of the training model constructed in step 2 to train the model automatically constructed by the neighbor graph index; during the training process, the loss function used to measure the error between the predicted performance and the actual performance The expression is:

[0035]

[0036] Among them, x i Automatically construct the input of the model for neighbor graph indexing; i is the actual performance; f(x i ) is the performance of the model prediction automatically constructed by the neighbor graph index; n is the number of groups of parameter configurations in the input model.

[0037] The loss function is calculated based on the predicted performance and the actual performance, and the model parameters of the feature extractor and parameter performance prediction module are adjusted through back propagation; the above training steps are repeated until the number of training iterations reaches the set number of iterations.

[0038] Step 5: Figure 2 As shown, the sample data set corresponding to the query data set is vectorized to obtain a vector data set; the vector data set is used to build a k-nearest neighbor graph, and multiple groups of parameter configurations are enumerated; the vector data set, the k-nearest neighbor graph, and the parameter configuration are input into the trained neighbor graph index to automatically build a model to obtain the prediction performance corresponding to the multiple groups of parameter configurations; the parameter configuration with the best prediction performance is selected to generate a neighbor graph index; after the query data set is vectorized, a neighbor search is performed on the generated neighbor graph index, thereby returning a certain number of results and completing the nearest neighbor search.

Claims

1. A nearest neighbor retrieval method based on automatic construction of a nearest neighbor graph index, characterized in that: The following steps are involved: Step 1: Obtain multiple sample data sets, perform feature extraction on the multiple sample data sets respectively, and obtain multiple vector data sets; the multiple vector data sets form a set V ; Step 2: Set V Each vector dataset in k Neighbor graph; for each k Each neighbor graph sets multiple sets of neighbor graph parameter configurations to build a neighbor graph index; obtain the performance indicators corresponding to each neighbor graph index; build a training set, which includes multiple samples, each sample includes a vector data set, k Neighborhood graph and its parameter configuration and performance indicators; Step 3: Establish a neighbor graph index automatic construction model; the neighbor graph index automatic construction model includes a feature extractor module and a parameter performance prediction module; the feature extractor module extracts and fuses the vector data set and k The features of the neighbor graph are used as the output of the feature extractor module; the input of the parameter performance prediction module is the concatenation of the output of the feature extractor module and the parameter configuration; the output of the parameter performance prediction module is the predicted performance index; Step 4: Use the training set constructed in step 2 to train the neighbor graph index automatic construction model; Step 5: Automatically build a model based on the trained neighbor graph index to obtain the optimal prediction performance index of the sample data set corresponding to the query data set; select the parameter configuration corresponding to the optimal prediction performance index to generate the neighbor graph index; after vectorizing the query data set, perform neighbor retrieval on the generated neighbor graph index.

2. According to claim 1, a nearest neighbor search method based on automatic construction of a nearest neighbor graph index is characterized in that: The feature extractor module includes a self-attention mechanism feature extraction module and a graph convolutional neural network; the self-attention mechanism feature extraction module is used to extract features from the vector data set; the graph convolutional neural network is used to extract features from the vector data set; k The nearest neighbor graph is used for feature extraction.

3. The nearest neighbor search method based on automatic construction of a nearest neighbor graph index according to claim 1, characterized in that: The parameter performance prediction module adopts a multi-layer perceptron.

4. The nearest neighbor retrieval method based on automatic construction of a nearest neighbor graph index according to claim 1, characterized in that: The sample data set is an image data set or a text data set; the images in the image data set are feature extracted using a convolutional neural network to obtain image vectors; the texts in the text data set are extracted using a pre-trained language model to obtain text vectors.

5. The nearest neighbor search method based on automatic construction of a nearest neighbor graph index according to claim 1, characterized in that: The performance metrics described include recall, number of queries processed per second, and number of distance calculations.

6. The nearest neighbor search method based on automatic construction of a nearest neighbor graph index according to claim 1, characterized in that: The neighbor graph parameter configuration includes the neighbor graph index type, the maximum number of neighbors, the number of nearest neighbors, and the size of the candidate set when composing the graph.

7. The nearest neighbor search method based on automatic construction of a nearest neighbor graph index according to claim 6, characterized in that: The method for setting multiple groups of parameter configurations is: enumerating different combinations of neighbor graph index type, maximum number of neighbors, number of nearest neighbors, and size of candidate set when composing a graph.

8. The nearest neighbor search method based on automatic construction of a nearest neighbor graph index according to claim 1, characterized in that: Step 4: Neighborhood graph index automatically builds the loss function of the model loss ( x , y ) is: ； in, x i Automatically construct model inputs for neighbor graph indexing; y i For real performance; f ( x i ) Automatically build model prediction performance for neighbor graph indexing; n The number of groups configured for the parameters in the input model.

9. A computer device comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein: The memory stores a computer program; the processor executes a nearest neighbor retrieval method automatically constructed based on a neighbor graph index as described in any one of claims 1-8.

10. A readable storage medium storing a computer program; characterized in that: When the computer program is executed by a processor, it is used to implement a nearest neighbor retrieval method based on automatic construction of a nearest neighbor graph index as described in any one of claims 1 to 8.

Images