Knowledge distillation method, electronic equipment and computer readable storage medium

Abstract

The invention discloses a knowledge distillation method, electronic equipment and a computer readable storage medium, and relates to the technical field of artificial intelligence, and the method comprises the steps: carrying out the grouping of an attention map matrix of a teacher network based on the number of attention heads in an initial network structure of a student network, a high-dimensional attention space of a teacher network can be divided into subspaces with the same dimension as a student network, and attention map matrixes in groups are spliced, so that each group of spliced matrixes is aligned with an attention head of an initial network structure, a relatively regular corresponding relationship between the attention maps of the teacher network and the student network is realized, and the attention map splicing efficiency is improved. Therefore, the distillation loss can be calculated, and the knowledge of the complex teacher network can be accurately transmitted to the lightweight student network. Therefore, the problem that in a multi-head attention mechanism, the number difference between the teacher network and the student network hinders the student network to fully learn the teacher network knowledge can be solved, and the technical effects of improving the knowledge distillation effect and enhancing the student network performance are achieved.

Description

Technical Field

[0001] The present application relates to the field of artificial intelligence technology, and in particular to a knowledge distillation method, an electronic device, and a computer-readable storage medium. Background Art

[0002] In today's rapidly developing artificial intelligence landscape, Transformer-based neural network models, with their superior feature extraction and processing capabilities, have become a core technology for intelligent transformation across various fields. However, Transformer models typically require significant storage and computing resources, limiting their application in resource-constrained scenarios.

[0003] To address these issues, related technologies use knowledge distillation to transfer the knowledge of a complex teacher network to a lightweight student network, reducing overhead while maintaining performance. However, when applying knowledge distillation to the Transformer model, the difference in the number of heads between the teacher and student networks in the multi-head attention mechanism makes it difficult to construct the distillation loss, making it difficult to accurately transfer the knowledge of the teacher network to the student network, hindering the student network from fully learning the teacher network's knowledge. Summary of the Invention

[0004] The present application provides a knowledge distillation method, an electronic device, and a computer-readable storage medium to at least solve the problem in the related art that the difference in the number of heads between the teacher network and the student network in the multi-head attention mechanism hinders the student network from fully learning the teacher network's knowledge.

[0005] This application provides a knowledge distillation method, including: Construct the initial network structure of the student network according to the network structure of the teacher network; Grouping multiple first attention map matrices of the teacher network; the number of the groups is the same as the number of attention heads in the initial network structure; Splicing the first attention map matrices in each group to obtain a splicing matrix corresponding to each group; The concatenated matrix and the second attention map matrix of the initial network structure are calculated according to the distillation loss function, and the network parameters of the initial network structure are updated based on the calculation results to obtain the student network.

[0006] The present application also provides an electronic device, comprising: a memory for storing a computer program; and a processor for implementing the steps of any of the above-mentioned knowledge distillation methods when executing the computer program.

[0007] The present application also provides a computer-readable storage medium, in which a computer program is stored, wherein when the computer program is executed by a processor, the steps of any of the above-mentioned knowledge distillation methods are implemented.

[0008] The present application also provides a computer program product, including a computer program, which implements the steps of any of the above-mentioned knowledge distillation methods when executed by a processor.

[0009] The present application also provides an image retrieval system, comprising an image encoding module and a text encoding module; The image encoding module is a student network obtained by any of the above-mentioned knowledge distillation methods, using the image encoder in the pre-trained model as the teacher network; The text encoding module is a student network obtained by any of the above-mentioned knowledge distillation methods using the text encoder in the pre-trained model as the teacher network.

[0010] Through this application, by grouping the attention graph matrix of the teacher network based on the number of attention heads in the initial network structure of the student network, the high-dimensional attention space of the teacher network can be divided into subspaces of the same dimension as the student network. By splicing the attention graph matrices within the group, each group of spliced ​​matrices is aligned with the attention heads of the initial network structure, achieving a more regular correspondence between the attention graphs of the teacher network and the student network, making the attention graph matrices of the teacher network and the student network comparable in form and structure, thereby being able to calculate the distillation loss and accurately transfer the knowledge of the complex teacher network to the lightweight student network. Therefore, the difference in the number of heads of the teacher network and the student network in the multi-head attention mechanism can be solved, which hinders the student network from fully learning the knowledge of the teacher network, thereby achieving the technical effect of improving the knowledge distillation effect and enhancing the performance of the student network. BRIEF DESCRIPTION OF THE DRAWINGS

[0011] In order to more clearly illustrate the embodiments of the present application, the following is a brief introduction to the drawings required for use in the embodiments. Obviously, the drawings described below are only some embodiments of the present application. For ordinary technicians in this field, other drawings can be obtained based on these drawings without any creative work.

[0012] Figure 1 This is a schematic diagram of the principle of the single-head attention mechanism provided in an embodiment of the present application; Figure 2 Schematic diagram of the process of the knowledge distillation method provided in the embodiment of the present application; Figure 3 Schematic diagram of the principle of the knowledge distillation method of the embodiment of the present application; Figure 4 Schematic diagram of the training principle of the image retrieval system according to an embodiment of the present application; Figure 5 Schematic diagram of the structure of the knowledge distillation device provided in the embodiment of the present application; Figure 6 It is a structural diagram of an electronic device provided in an embodiment of the present application. DETAILED DESCRIPTION

[0013] The following will be combined with the accompanying drawings in the embodiments of this application to clearly and completely describe the technical solutions in the embodiments of this application. Obviously, the embodiments described are only part of the embodiments of this application, not all of them. Based on the embodiments in this application, all other embodiments obtained by ordinary technicians in this field without making creative efforts are within the scope of protection of this application.

[0014] It should be noted that, in the description of this application, the terms "comprises," "includes," or any other variations thereof are intended to cover non-exclusive inclusion, such that a process, method, article, or device comprising a series of elements includes not only those elements, but also other elements not explicitly listed, or elements inherent to such process, method, article, or device. The terms "first," "second," etc., in this application are used to distinguish similar objects, and are not used to describe a particular order or sequence.

[0015] The core idea of ​​the attention mechanism is to calculate the output based on the query, key, and value. , bond matrix Sum Matrix ,in is the number of queries, keys and values, is the dimension of the query and key, is the dimension of the value. express OK, A real matrix with columns.

[0016] During the operation of the attention mechanism, the similarity between the query and the key is first calculated to obtain the attention score matrix S:

[0017] In the above formula, Superscript Represents matrix transpose. Represents a real matrix with n rows and n columns. In order to avoid the dot product result being too large and affecting the gradient stability during training, a scaling operation is usually performed, divided by , get the scaled attention score :

[0018] Then use the Softmax function to convert the scaled attention scores into probability distributions to obtain the attention map matrix A:

[0019] in, The function converts a row of elements into probability values.

[0020] Finally, the attention map matrix and value matrix Multiply them together to get the output of the attention mechanism :

[0021] Figure 1 Schematic diagram of a single-head self-attention mechanism. Figure 1 In the example, assuming that the current layer has input embedding vector, The input embedding vector is , . Input embedding vector Input to query value transformation matrix Get output query , input embedding vector Input to the key-value transformation matrix Get output key ; Input embedding vector Input to feature embedding value transformation matrix Get value . The query matrix is ​​expressed as , bond matrix , value matrix In the self-attention mechanism, the inputs of the three transformation matrices are the same, while in the non-self-attention mechanism, the inputs of the three transformation matrices are different.

[0022] The similarity calculation layer between the query and the key calculates the similarity between the query and the key to obtain the attention score matrix, and the attention score matrix is ​​scaled in the scaling layer to obtain the scaled attention score. Then, the Softmax function is used in the Softmax layer to convert the scaled attention score into a probability distribution to obtain the attention map matrix. Finally, in the output calculation layer, the attention map matrix is ​​multiplied by the value matrix to obtain the output of the attention mechanism, that is, the feature embedding vector .here, , , , , , , are all column vectors.

[0023] The multi-head attention mechanism calculates the attention map matrix in parallel through multiple different attention heads, enabling the model to capture information from different representation subspaces.

[0024] Assume there is Attention heads. For the input query matrix , key matrix Sum Matrix , respectively through A set of different transformation matrices , and , Perform the transformation, where is the dimension of the model, is the dimension of query and key, is the dimension of the value. .

[0025] No. The query matrix of the attention heads , key matrix Sum Matrix They are:

[0026]

[0027] For each attention head, the output is calculated following the steps of the basic single-head attention mechanism described above:

[0028]

[0029] in, Indicates the The attention map matrix of the attention heads, Superscript Representative Matrix Transpose, Indicates the The output of an attention head.

[0030] Will The output of an attention head Splice them together to get the spliced ​​matrix :

[0031] in, Indicates n rows, Real matrix of columns Finally, through a linear transformation matrix Transform the concatenated matrix to obtain the final output of the multi-head attention mechanism :

[0032] Knowledge distillation is a model compression technology that aims to transfer the knowledge contained in a complex and powerful teacher network to a lightweight student network. This allows the student network to maintain performance close to that of the teacher network while significantly reducing the number of parameters and computational complexity. It is particularly suitable for scenarios with limited resources.

[0033] In the knowledge distillation framework, the teacher network typically consists of a large number of parameters, possesses excellent feature extraction and processing capabilities, and is capable of learning complex knowledge representations. The teacher network is typically a fully trained deep neural network with a deep hierarchical structure, such as ResNet (Residual Network) and BERT (Bidirectional Encoder Representations from Transformers). Due to its large number of parameters and complex structure, it is able to learn rich feature information and advanced representations. The teacher network provides "guidance" to the student network by outputting soft labels or intermediate features, enabling the student network to better understand the underlying structure in the data.

[0034] The student network is a relatively simple and lightweight model, such as MobileNet (Mobile Network-based models) and TinyBERT (Tiny Bidirectional Encoder Representations from Transformers). Its goal is to mimic the behavior of the teacher network and efficiently complete tasks within limited computing resources. The student network has a small number of parameters and layers. By absorbing the knowledge embedded in the teacher network, it can achieve satisfactory accuracy while maintaining low computational cost. During training, the student network learns the soft labels or intermediate features output by the teacher network, gradually approaching the teacher network's performance, ultimately achieving model optimization and compression.

[0035] In the knowledge distillation process based on the Transformer architecture, the difference in the number of heads in the multi-head attention mechanism between the teacher network and the student network leads to an inconsistency in the number of attention map matrices between the two networks. This makes the traditional distillation loss based on a one-to-one comparison of attention map matrices difficult to directly apply. If forced to use this loss, the mismatch in the number of attention maps will cause information loss, seriously affecting the student network's effective learning and absorption of the teacher network's knowledge, and failing to achieve the goal of knowledge distillation guidance or improving the student network's performance.

[0036] To solve the above problems, an embodiment of the present application provides a knowledge distillation method.

[0037] In order to enable those skilled in the art to better understand the present application, the present application is further described in detail below with reference to the accompanying drawings and specific implementation methods.

[0038] In combination with the specific application environment architecture or specific hardware architecture on which the execution of the knowledge distillation method depends, the specific application environment architecture or specific hardware architecture is described here.

[0039] The embodiments of the present application provide a knowledge distillation method, and the method is described in detail in conjunction with the execution process of the knowledge distillation method.

[0040] Specifically, Figure 2 The present invention provides a flowchart of a knowledge distillation method according to an embodiment of the present application.

[0041] like Figure 2 As shown, the knowledge distillation method includes: step 210, step 220, step 230 and step 240.

[0042] Step 210: Construct an initial network structure of the student network based on the network structure of the teacher network.

[0043] The teacher network's architecture includes several factors, such as the number of layers, the number of neurons in each layer, the choice of activation function, and the number of attention heads in the multi-head attention mechanism. When constructing the initial structure of the student network, the teacher network's architecture can be analyzed and evaluated. For example, the teacher network may have deeper layers and a larger number of attention heads to enable complex feature extraction and information exchange. While maintaining a certain level of performance, the student network should reduce the number of layers and attention heads to lower computational complexity and memory usage.

[0044] When constructing the student network's initial network structure, the topology and connection logic of the teacher network's network structure can be retained. For example, if the teacher network adopts the basic block structure of "multi-head attention layer + feedforward network + residual connection", the student network's initial network structure can also follow the same basic block structure, ensuring that the feature processing flow of the initial network structure is consistent with that of the teacher network. To achieve lightweightness, the number of attention heads in the student network's initial network structure can be reduced, allowing the student network to inherit the capabilities of the teacher network while meeting lightweight requirements.

[0045] Step 220: Group multiple first attention map matrices of the teacher network; the number of groups is the same as the number of attention heads in the initial network structure.

[0046] In the embodiments of this application, based on the introduction to the multi-head attention mechanism described above, attention heads are independent computational units within the multi-head attention mechanism. Each head captures the features and relationships of the input data from a different perspective, improving the model's understanding of information through parallel computation. The attention map matrix records the distribution of attention scores between various positions during the attention mechanism's computation. Each attention head generates a corresponding attention map matrix, which together constitute the multi-head attention mechanism's comprehensive understanding of the input data.

[0047] In an embodiment of the present application, the attention map matrix of the teacher network is called the first attention map matrix. In the knowledge distillation scenario based on the Transformer architecture, the number of attention heads in the teacher network is greater than the number of attention heads in the student network, and the number of attention map matrices generated by the teacher network is also greater than the number of attention map matrices generated by the student network. The attention heads of the initial network structure can be grouped so that the number of groups is consistent with the number of attention heads in the initial network structure. Such a grouping operation reorganizes the high-dimensional attention information of the teacher network and converts it into a form that is compatible with the student network structure, laying the foundation for information alignment between the teacher network and the student network in the subsequent knowledge distillation process.

[0048] For example, the number of attention heads of the initial network structure is 4, and the number of attention heads of the teacher network is 12. Correspondingly, the number of attention map matrices of the teacher network is also 12. According to the number of attention heads of the initial network structure, the 12 attention map matrices of the teacher network can be divided into 4 groups, so that the teacher network structure is adapted to the student network structure, that is, one attention head of the student network corresponds to a group of attention map matrices of the teacher network.

[0049] Step 230: Splice the first attention map matrices in each group to obtain a spliced ​​matrix corresponding to each group.

[0050] Since each group contains multiple first attention map matrices of the teacher network, although these attention map matrices contain rich attention information, they are independent and dimensionally dispersed, making it difficult to directly compare and transfer knowledge with the attention map matrix of the student network.

[0051] Therefore, the first attention map matrices in each group can be spliced ​​to obtain the corresponding splicing matrices of each group. The splicing can be done by column splicing, element-wise weighted splicing, or other splicing methods, which are not limited in this embodiment of the application.

[0052] In an embodiment of the present application, through matrix splicing, the multiple attention map matrices originally scattered in each group are merged into a splicing matrix with a specific dimension, which not only compresses the dimension of the teacher network attention information, making the splicing matrix more consistent with the dimension of the attention map matrix corresponding to the single attention head of the initial network structure, but also integrates the attention information from different angles and levels in the same group. In this way, each splicing matrix is ​​equivalent to condensing the information of multiple attention heads of the teacher network into a feature representation that can directly correspond to a single attention head of the initial network structure, so that the re-arranged attention knowledge of the teacher network can be presented in a more regular and comparable form.

[0053] In some embodiments, the dimension of the splicing matrix is ​​the same as the dimension of the attention map matrix of the initial network structure, so that the splicing matrix can be compared one-to-one with the dimension of the attention map matrix corresponding to a single attention head of the initial network structure, so that the attention map matrices of the teacher and the initial network structure are comparable in form and structure, thereby being able to calculate the distillation loss and accurately transfer the knowledge of the complex teacher network to the lightweight student network.

[0054] Step 240: Calculate the concatenated matrix and the second attention map matrix of the initial network structure according to the distillation loss, and update the network parameters of the initial network structure based on the calculation results to obtain the student network.

[0055] In the embodiment of the present application, the distillation loss is a quantitative indicator to measure the knowledge difference between the teacher network and the student network. Its function is to guide the student network to learn the knowledge of the teacher network so that the output of the student network is as close as possible to the teacher network.

[0056] In the knowledge distillation scenario of the Transformer model, due to the difference in the number of heads in the multi-head attention mechanisms of the teacher and student networks, directly calculating the loss function between the teacher and student attention map matrices can hinder accurate transfer of the teacher's knowledge to the student due to dimensionality mismatch and structural inconsistency. However, after grouping and concatenating the teacher's attention map matrices, the concatenated matrix of the teacher and the second attention map matrix of the initial network structure align in dimensionality and structure, making them suitable for direct comparison. Based on this, distillation loss functions such as KL (Kullback-Leibler) divergence, mean squared error (MSE), and cross-entropy loss can be used to construct a measure of the difference between the concatenated matrix and the second attention map matrix of the initial network structure. For example, using KL divergence to construct a distillation loss can capture semantic deviations in the attention weight distribution patterns between the concatenated matrix of the teacher and the student's attention map by calculating the difference in probability distribution between them.

[0057] By calculating the difference between the concatenated matrix and the second attention map matrix of the initial network structure, the knowledge such as attention distribution and feature association contained in the teacher network is quantified in the form of a distillation loss value. This loss value intuitively reflects the gap between the current state of the initial network structure and the expected state of the teacher network.

[0058] In an embodiment of the present application, during the training process, the constructed distillation loss will be used as the optimization target. The initial network structure adjusts its own parameters according to the size of the distillation loss through the back-propagation algorithm, continuously narrowing the gap in attention features with the teacher network, and gradually learning the attention mechanism knowledge of the teacher network, thereby achieving effective migration of the knowledge of the complex teacher network to the lightweight student network, while reducing the model resource overhead and ensuring the performance of the student network.

[0059] Specifically, during training, the training dataset can be fed simultaneously into the teacher network and the initial network structure. The training objectives consist of two parts: first, measuring the difference between the initial network structure and the teacher network's attention map through distillation loss; second, ensuring the initial network structure's ability to learn the original task through task loss. The distillation loss and task loss can be weighted together to form a total loss, which serves as the basis for guiding the update of the initial network structure parameters.

[0060] In practical applications, the student network can be applied to different scenarios to achieve different tasks. For example, in the field of computer vision, visual models pre-trained on large-scale image datasets, such as ResNet (Residual Network) and EfficientNet (Efficient Convolutional Neural Network), can be used as teacher networks. Student networks trained through knowledge distillation can be applied to tasks such as image classification, object detection, and image retrieval. For example, in medical image diagnosis, the student network can distinguish between normal and diseased tissue; in security monitoring, the student network can identify objects such as people and vehicles in surveillance footage; and in autonomous driving applications, the student network can also identify pedestrians and traffic signs on the road through object detection.

[0061] In the field of natural language processing, pre-trained language models such as BERT (Bidirectional Encoder Representations from Transformers) and GPT (Generative Pre-trained Transformer) can be used as teacher networks. Student networks trained through knowledge distillation can then be deployed for semantic understanding, text classification, sentiment analysis, or question-answering systems. For example, in text classification tasks, the student network can categorize news articles by topic, helping news platforms organize and recommend relevant content. In machine translation tasks, the student network can understand the semantics of the source language text and translate it into the target language.

[0062] In multimodal interaction scenarios, pre-trained multimodal models such as CLIP (Contrastive Language-Image Pre-training) and ALBEF (Aligning Bidirectional Encoder Representations from Transformers with Fine-tuning) can be used as teacher networks. Student networks trained through knowledge distillation can then be applied to multimodal interaction tasks. For example, in the image-based question-and-answer function of intelligent assistants, the student network can process user-entered image and text questions, combining visual and textual information to generate answers. In product search scenarios on e-commerce platforms, the student network can accurately match matching products from a product image library based on user-entered text descriptions.

[0063] By distilling the attention mechanism learned by the teacher network, the student network can inherit the capabilities of the teacher network while reducing the computational burden and is suitable for running on mobile or edge devices.

[0064] The training datasets used by the student network vary depending on the application scenario. For natural language tasks, text datasets can be used. These can include multilingual corpora, annotated sentiment analysis datasets, question-and-answer datasets, news articles, or domain-specific corpora (such as medical or legal texts). This helps the student network learn deep semantics and attention allocation mechanisms. For image processing tasks, image datasets such as ImageNet, detection datasets with annotated object locations and categories, such as the COCO (Common Objects in Context) dataset, or continuous video frame sequences can be used.

[0065] The goal of knowledge distillation is to make the initial network structure lightweight while maintaining good task performance. For example, in an image classification task, the initial network structure needs to correctly identify the object category in the image. The task loss calculates the difference between the predicted category probability distribution of the initial network structure and the true label, allowing the initial network structure to imitate the knowledge of the teacher network while also achieving the actual task objectives. Task loss can be cross entropy loss, mean squared error, etc.

[0066] During training, the backpropagation algorithm propagates the total loss to each parameter layer of the initial network structure, calculating the gradient of each parameter and adjusting the weights accordingly. For example, when the distillation loss indicates that the output of a particular attention head in the initial network structure differs significantly from the concatenated matrix of the teacher network, the gradient update will focus on adjusting the weight parameter corresponding to that attention head, aligning its attention distribution with that of the teacher network. Conversely, when the task loss is high, it indicates that the initial network structure's predictions for actual tasks (such as classification and regression) deviate significantly from the true labels. The backpropagation algorithm will then prioritize adjusting parameters that affect the task output. For example, in an image classification task, if the initial network structure mistakenly classifies a cat as a dog, the task loss will calculate the gradient based on the difference between the prediction and the true label. This in turn significantly updates the parameters of the convolutional and fully connected layers responsible for feature extraction and final decision-making, enabling the initial network structure to more accurately distinguish between different categories in subsequent predictions. By synergizing the distillation loss with the task loss, the initial network structure can optimize the characteristics of the attention mechanism while continuously strengthening its ability to perform actual tasks.

[0067] Through this application, by grouping the attention graph matrix of the teacher network based on the number of attention heads in the initial network structure of the student network, the high-dimensional attention space of the teacher network can be divided into subspaces of the same dimension as the student network. By splicing the attention graph matrices within the group, each group of spliced ​​matrices is aligned with the attention heads of the initial network structure, achieving a more regular correspondence between the attention graphs of the teacher network and the student network, making the attention graph matrices of the teacher network and the student network comparable in form and structure, so that the distillation loss can be calculated, and the knowledge of the complex teacher network can be accurately transferred to the lightweight student network. Therefore, the difference in the number of heads of the teacher network and the student network in the multi-head attention mechanism can be solved, which hinders the student network from fully learning the knowledge of the teacher network, thereby achieving the technical effect of improving the knowledge distillation effect and enhancing the performance of the student network.

[0068] In some embodiments, grouping a plurality of first attention map matrices of a teacher network includes: Calculate the ratio of the number of attention heads in the teacher network to the number of attention heads in the student network; Determine a grouping strategy based on the ratio; Multiple first attention map matrices of the teacher network are grouped according to the grouping strategy.

[0069] In this embodiment, the architectural difference between the teacher network and the student network in the multi-head attention mechanism can be quantified by the number of attention heads in the teacher network and the initial network structure.

[0070] Taking a layer corresponding to the teacher network and the initial network structure as an example, the multi-head attention mechanism of the teacher network in this layer includes attention heads, and correspondingly generate First attention map matrix , ,……, ; In the initial network structure of the same layer, the number of attention heads of the multi-head attention mechanism is , and accordingly produces The second attention map matrix is , ,……, ,generally, , the ratio can be calculated .

[0071] After obtaining the ratio, we can dynamically design a grouping strategy based on the ratio. For example, when the ratio is an integer, we can directly use equal number grouping, that is, every 3 heads in the teacher network corresponds to 1 head in the student network. If the ratio is non-integer, we can introduce a mixed grouping strategy: some groups contain 3 heads, and some groups contain 2 heads.

[0072] In some embodiments, during the grouping process, consecutive first attention map matrices in the teacher network can be grouped together. For example, the 12 first attention map matrices of the teacher network are numbered sequentially, and each group of three consecutive first attention map matrices is grouped together, forming four groups. Through this continuous grouping approach, the first attention map matrices within each group maintain the integrity of the local features of the teacher network.

[0073] In some embodiments, the grouping strategy can also be combined with the semantic relevance of the attention map matrix, and attention map matrices with complementary features are preferentially grouped into the same group. For example, if one attention map matrix focuses on local syntax and the other attention map matrix focuses on global semantics, it means that the two can complement each other.

[0074] Specifically, we can calculate the similarity between the first attention map matrices, such as cosine similarity and Euclidean distance. The higher the similarity, the higher the semantic relevance of the two first attention map matrices and the more similar the information they focus on. The lower the similarity, the lower the semantic relevance of the two first attention map matrices and the more different the information they focus on, which can complement each other. Therefore, we can group the first attention map matrices with low similarity into one group.

[0075] In this embodiment, by calculating the similarity between the attention map matrices of the teacher network, the attention map matrices with low similarity are preferentially classified into the same group, so that each group includes heterogeneous features. This feature complementarity enables each group to capture fine-grained grammatical structures and macro-semantic relationships, thereby improving the richness of knowledge representation; and by reducing the redundant combinations of high-similarity attention map matrices, the dilution of the distillation signal caused by repeated information within the group is reduced, so that the subsequent distillation loss calculation can more accurately focus on the key feature differences.

[0076] In this embodiment, by calculating the ratio of the number of attention heads in the teacher network to the number of attention heads in the initial network structure, and dynamically grouping the ratio according to the strategy, adaptive compression of the teacher network's attention knowledge and structured alignment with the student network are achieved, thereby enabling the construction of distillation loss and accurate transfer of the knowledge of the complex teacher network to the lightweight student network.

[0077] In some embodiments, the plurality of first attention map matrices of the teacher network are grouped according to a grouping strategy, comprising: When the ratio is an integer, the grouping strategy includes: dividing the multiple first attention map matrices equally into group; among them, Equal to the number of attention heads in the student network.

[0078] In this embodiment, when the ratio of the number of attention heads of the teacher network to the number of attention heads of the student network is calculated to be an integer, the number of attention heads of the teacher network can be The first attention map matrix is ​​evenly divided into groups, each containing exactly In this way, a more regular correspondence between the attention maps of the teacher network and the student network can be achieved.

[0079] For example, if the teacher network has 12 attention heads and the student network has 4 attention heads, with a ratio of 3, an equal number division strategy can be adopted. The 12 first attention map matrices of the teacher network are numbered sequentially, and every 3 consecutive first attention map matrices are divided into a group, forming 4 groups, which has the same number of attention heads as the student network. Each group can correspond to an attention head of the student network, so that each group of first attention map matrices can generate a distillation target aligned with the student network through simple splicing or weighted averaging, without the need for additional dimensionality transformation or complex mapping.

[0080] In this embodiment, through the equal number grouping strategy under the integer ratio condition, the knowledge mapping from the multi-head teacher to the small-head student is completed with simple logic, so that the grouped matrix is ​​comparable with the student network in terms of structural distribution, thereby enabling the construction of distillation loss and achieving the accurate transfer of knowledge from the complex teacher network to the lightweight student network.

[0081] In some embodiments, grouping a plurality of first attention map matrices of a teacher network includes: When the ratio is not an integer, multiple first attention map matrices are divided into group; of which, at most -1 The number of first attention maps in the group is the ratio rounded down; Equal to the number of attention heads in the student network.

[0082] In this embodiment, the ratio is not an integer, i.e. no The integral multiple of the teacher network can be The first attention map matrix is ​​also divided into Group. Among them, -1 group each contains feature maps, and the last group contains feature maps. ,in Indicates floor operation.

[0083] In this embodiment, a hybrid grouping strategy for non-integer ratios achieves the unification of elastic compression and dimensional alignment. Specifically, the non-divisibility grouping problem is solved through a mechanism of floor rounding and remainder distribution, ensuring consistency between the number of groups and the number of attention heads in the student network. This enables the construction of a distillation loss and the accurate transfer of knowledge from a complex teacher network to a lightweight student network.

[0084] In some embodiments, calculating the concatenation matrix and the second attention map matrix of the initial network structure according to the distillation loss includes: Reduce the dimension of the concatenated matrix so that the dimension of the concatenated matrix is ​​the same as the dimension of the second attention map matrix; The calculation results are obtained based on the concatenated matrix after dimensionality reduction and the second attention map matrix.

[0085] In this embodiment, after the teacher network attention map matrix is ​​grouped and concatenated, although the concatenated matrix corresponds to the attention head of the initial network structure in structure, due to the architectural differences between the teacher network and the initial network structure, the concatenated matrix and the second attention map matrix of the initial network structure may still not match in dimension. For example, the concatenated matrix of the teacher network may differ from the second attention map matrix of the student network in terms of the number of channels, sequence length, or feature dimension. For example, the dimension of the concatenated matrix is ​​128×128, while the dimension of the second attention map matrix is ​​64×64, making it difficult to directly construct the distillation loss.

[0086] In this embodiment, a dimensionality reduction operation may be performed on the splicing matrix so that the dimension of the splicing matrix after dimensionality reduction is consistent with the dimension of the second attention map matrix.

[0087] The core goal of dimensionality reduction is to adjust the dimension of the splicing matrix through specific mathematical transformations while retaining key feature information, so that the dimension of the splicing matrix after dimensionality reduction is consistent with the dimension of the second attention map matrix of the student network. In this embodiment, the dimensionality reduction method adopted includes: using pooling operations (such as average pooling and maximum pooling) to compress the spatial dimension of the splicing matrix; or mapping and adjusting the feature dimension of the splicing matrix through neural network layers such as fully connected layers and convolutional layers. For example, if the sequence length of the splicing matrix is ​​greater than the second attention map matrix of the student network, the sequence length can be reduced by average pooling; if the feature dimensions of the two are different, the feature dimensions of the splicing matrix can be mapped to the dimension consistent with the student network using a fully connected layer.

[0088] In this embodiment, the splicing matrix of the teacher network is adjusted through a dimensionality reduction operation so that the dimension of the splicing matrix after dimensionality reduction is the same as the dimension of the second attention map matrix of the initial network structure, so that they can be compared in a unified dimensional space. The difference between the teacher network and the initial network structure can then be accurately quantified through the distillation loss, and the attention knowledge contained in the teacher network is fed back to the initial network structure in the form of a loss value, further improving the efficiency and accuracy of knowledge distillation.

[0089] In some embodiments, performing dimensionality reduction on the concatenated matrix includes: The concatenated matrix is ​​reduced in dimensionality using a target matrix, where the target matrix is ​​a learnable matrix that is updated by the back-propagation algorithm during the student network training process.

[0090] In this embodiment, before training starts, a learnable target matrix can be initialized for the splicing matrix. The target matrix dimensions are designed to be the number of rows in the splicing matrix × the number of columns in the student attention map matrix (e.g., 128 × 64). During training, real-time dimensionality reduction can be achieved through matrix multiplication: the splicing matrix (e.g., 128 × 128 dimensions) is multiplied by the target matrix (e.g., 128 × 64 dimensions) to obtain a dimensionally aligned output matrix (128 × 64).

[0091] For example, for the Group( ), by First Attention Map , ,……, Splice by column to form a splicing matrix , the concatenation matrix dimension is , represent OK, The second attention map matrix of the initial network structure is , .because, and The dimensions of are different and the distillation loss cannot be constructed directly, so a learnable target matrix is ​​introduced , the dimension of the target matrix is , through matrix multiplication operation ,Will The dimension is transformed to The dimensions are the same.

[0092] After each forward computation, the distillation loss between the concatenated matrix after dimensionality reduction and the second attention map matrix of the initial network structure is calculated, and the gradient is transferred to the target matrix via backpropagation. The update rule can use an adaptive learning rate optimizer, such as Adam, to dynamically adjust the weights of the target matrix based on the gradient direction. It is important to note that the target matrix is ​​updated simultaneously with the network parameter updates of the initial network structure. On the one hand, the target matrix learns the optimal projection direction, and on the other hand, the student network learns how to match the feature distribution after projection, forming a closed loop between the dimensionality reduction process and knowledge transfer.

[0093] Figure 3 Schematic diagram of the principle of the knowledge distillation method of an embodiment of the present application. Figure 3The following figure shows the attention map matrix of a layer of the teacher network and the corresponding attention map matrix of a layer of the initial network structure. First, the teacher network's attention map matrices are grouped according to the number of attention heads in the initial network structure. Then, the concatenation of each group of attention map matrices is reduced to the same dimension as the initial network structure's attention map matrix using the target matrix. Finally, the distillation loss is calculated using the reduced concatenation matrix and the initial network structure's attention map matrix.

[0094] In this embodiment, by adopting a learnable matrix that can be updated with the propagation algorithm as the target matrix for dimensionality reduction, the projection direction can be continuously optimized. Compared with the fixed dimensionality reduction method, more key knowledge of the teacher network can be retained, and the dimensionality reduction process can be dynamically adjusted according to the learning progress of the initial network structure, further improving the efficiency of knowledge transfer.

[0095] In some embodiments, calculations are performed based on the concatenated matrix after dimensionality reduction and the second attention map matrix, including: The target diagonal matrix is ​​used for constraint, and the concatenated matrix after dimensionality reduction and the second attention map matrix are calculated based on the objective function to obtain the calculation result; the objective function represents the difference between the teacher network and the initial network structure.

[0096] In this embodiment, after completing the dimensionality reduction of the concatenated matrix, it is necessary to construct an objective function that accurately measures the difference between the teacher network and the student network. This objective function can be constructed using a variety of distance metrics, such as mean square error (MSE) and KL (Kullback-Leibler) divergence.

[0097] To further enhance the distillation effect, a target diagonal matrix is ​​introduced as a constraint when constructing the distillation loss. A target diagonal matrix is ​​a matrix where the diagonal elements are learnable and the remaining elements are zero. The target diagonal matrix is ​​used to weight the objective function, ensuring that features of different dimensions have different importance in the knowledge transfer process.

[0098] In some embodiments, when yes When the integer multiple of Group( ) can be defined as:

[0099] in, Indicates the Layer, The distillation loss of the group, represents the objective function, Indicates the The concatenation matrix of the group, represents the target matrix, represents the second attention map matrix, and represents different target diagonal matrices, represents the hyperparameter, It means taking the main diagonal elements of the matrix to form a column vector, Represents the zero norm of the matrix, which enables the target diagonal matrix to be selectively activated and effective.

[0100] The formula for KL divergence is:

[0101] Among them, P is the attention distribution of the teacher network, and Q is the attention distribution of the initial network structure.

[0102] exist When it is KL divergence:

[0103] in, Indicates that each row element value in the matrix is ​​converted into a probability value.

[0104] In some embodiments, when no When the teacher network is an integer multiple of Groups, each containing The first attention map matrix, according to yes In the same way as the integers, the attention maps in each group are concatenated into a new matrix ( ), and adjust the dimension through the target matrix, and then construct the distillation loss in combination with the target diagonal matrix.

[0105] And for the last group, including The first attention map is recorded as , ,……, . Similarly, these attention maps are first spliced ​​by column to form a splicing matrix ,at this time Since its dimension corresponds to the feature map of the initial network structure Inconsistent, still using the target matrix right Perform dimension transformation. Then, introduce the target diagonal matrix and .use Constraints are set to more effectively screen out the most valuable feature columns for student network learning while retaining key information, and finally construct the distillation loss function:

[0106] in, Indicates the Layer, distillation loss of the group; represents the objective function, which can be KL divergence; represents the first The group first attention map matrix concatenation matrix, is the target matrix; and are two target diagonal matrices, in which the values ​​of elements other than the main diagonal elements are zero; It means taking the main diagonal elements of the matrix to form a column vector, Represents the zero norm of the matrix, which is used to enable the target diagonal matrix to be selectively activated. is a hyperparameter.

[0107] The above processing method is for a certain layer of the teacher network and the initial network structure, and it is also applicable to other layers in the Transformer architecture. This is because the multi-head attention mechanism structure of each layer in the Transformer architecture is similar. Although the input and output features of different layers may differ, the difference in the number of heads in the multi-head attention mechanism causes the same problem of mismatch in the attention map matrix.

[0108] In this embodiment, the objective function can accurately capture the difference between the teacher network and the initial network structure, providing a clear optimization goal for knowledge transfer. The introduction of the target diagonal matrix realizes differentiated weighting of different feature dimensions and can selectively activate effective columns, so that the training process can focus on key knowledge, thereby improving the pertinence and efficiency of knowledge transfer.

[0109] In some embodiments, the target diagonal matrix includes multiple diagonal matrices that are updated through the back-propagation algorithm during the initial network structure training process.

[0110] In this embodiment, before distillation training begins, multiple independent diagonal matrices can be initialized. The dimensions of the diagonal matrices are the same as those of the concatenation matrix and the second attention map matrix. These diagonal matrices can be initialized to the identity matrix, with initial weights assigned to the main diagonal elements. Each diagonal matrix is ​​marked as a learnable parameter, allowing it to be updated via backpropagation during training.

[0111] Specifically, when the initial network structure begins training, the backpropagation algorithm starts as each round of forward propagation calculates the distillation loss based on the concatenated matrix after dimensionality reduction and the second attention map matrix. At this time, the gradient is not only transmitted and updated along the parameter path of the initial network structure, but also transmitted to these diagonal matrices. Since the diagonal elements of the diagonal matrix correspond to the weights of different feature dimensions, the gradient calculated by backpropagation will adjust the weight of each diagonal element according to the network differences reflected by the distillation loss. For example, if the feature dimension corresponding to a diagonal matrix contributes less to narrowing the network differences in knowledge transfer, the gradient of the diagonal element will guide the weight to decrease, thereby achieving dynamic adaptive updates of the diagonal matrix.

[0112] In this embodiment, by designing the target diagonal matrix as multiple updateable diagonal matrices, the importance of knowledge of different dimensions in the teacher network can be adaptively explored, and the initial network structure can be guided to learn key knowledge. The synergistic effect of multiple diagonal matrices can more comprehensively and accurately measure the difference between the teacher network and the initial network structure, so that the initial network structure can not only learn the overall knowledge distribution of the teacher network during training, but also deeply absorb the core features of each dimension, thereby improving the efficiency and quality of knowledge distillation.

[0113] In some embodiments, the concatenation matrix and the second attention map matrix of the initial network structure are calculated according to the distillation loss, and the network parameters of the initial network structure are updated based on the calculation results to obtain a student network, including: Get a labeled training dataset; The network parameters of the teacher network are fixed, the training dataset is input into the teacher network and the initial network structure, and the network parameters of the initial network structure are updated through the back propagation algorithm based on the overall loss function; the overall loss function is constructed based on the distillation loss and the task loss; the task loss represents the difference between the predicted value of the initial network structure and the true value corresponding to the label.

[0114] In this embodiment, the training dataset serves as the basis for initial network structure learning. The training dataset may include input samples and their corresponding true labels. For example, in an image classification task, the training dataset consists of image samples, and the labels are the categories to which the images belong (e.g., cat, dog, etc.).

[0115] The overall loss function consists of distillation loss and task loss. Distillation loss measures the difference in attention mechanisms between the initial network structure and the teacher network. It is constructed through operations such as grouping, concatenating, and dimensionality reduction of the attention map matrix described above, allowing the initial network structure to mimic the feature extraction methods of the teacher network. Task loss, on the other hand, is used to evaluate the predictive ability of the student network and is quantified by calculating the difference between the output of the initial network structure and the true label. Task loss can be implemented as cross-entropy loss, mean squared error, or other methods.

[0116] Taking the task of the initial network structure as an example, which is to predict the next token given a sequence of tokens, the task loss function is A cross entropy loss function can be used.

[0117] Suppose there is a A sequence of word units, for each position , the student network needs to predict the next word in the vocabulary The probability distribution of words. is a A one-hot vector of dimension denoting the position The category corresponding to the real word at It is predicted by the student network dimensional probability distribution vector, then the cross entropy loss function can be expressed as:

[0118] in, Indicates mission loss, Indicates the total number of words, represents the vocabulary size, yes The vector elements, The vector is dimension hot vector, indicating the true corresponding category at position i, yes The vector elements, The initial network structure predicts a V-dimensional probability distribution vector, representing the probability that the initial network structure predicts that the token at position i belongs to each word in the vocabulary. This formula calculates the average difference between the probability distribution predicted by the initial network structure and the true label. By minimizing this loss function, the parameters of the student network are adjusted, allowing the initial network structure to better predict the next token, thereby improving the performance of the language model.

[0119] In this embodiment, by combining distillation loss and task loss, the dual goals of knowledge transfer and task optimization are achieved. The distillation loss enables the student network to continuously absorb the high-order attention patterns of the teacher network during the parameter update process, and the supervision of the task loss enables the transferred knowledge to serve the real task goals, thereby improving the student network's ability to handle actual tasks.

[0120] In some embodiments, the overall loss function is constructed as follows: Construct an overall distillation loss based on the distillation loss; the overall distillation loss represents the overall difference between the teacher network and the different network layers of the initial network structure; The overall loss function is constructed based on the overall distillation loss and the task loss.

[0121] In this example, the distillation loss of the above construction Indicates the Layer, The distillation loss of the group, Indicates the Layer, The distillation loss of the group. The teacher network and the student network are both L layers, and can be calculated based on the distillation loss and distillation losses Construct the overall distillation loss.

[0122] Specifically, assuming that the teacher network The number of attention heads in the layer is , the attention map matrix is , ,…, ; Initial network structure The number of attention heads in the layer is , the attention map matrix is , ,…, ,in For the layer, calculation , and follow the above rules to manage the teacher network The attention map matrix is ​​grouped.

[0123] When the ratio of the number of attention heads in the teacher network to the number of attention heads in the initial network structure is an integer, the overall distillation loss function is Defined as:

[0124] in, Representative Layer, The distillation loss of the group.

[0125] The above expression can be expanded into the following formula:

[0126] in, is the objective function, which can be KL divergence; It is the first Layer A concatenation matrix consisting of a set of first attention map matrices, It is The target matrix of the layer is used to convert The dimension is converted to The dimensions are the same, It is the first Layer A second attention map matrix; and is the target diagonal matrix; It means taking the main diagonal elements of the matrix to form a column vector.

[0127] When the ratio of the number of attention heads in the teacher network to the number of attention heads in the initial network structure is not an integer, the overall distillation loss function Defined as:

[0128] Because the last group requires special treatment, the above formula can also be expanded into:

[0129] More specifically, when When:

[0130] in, Representative Layer, The distillation loss of the group, is the objective function, which can be KL divergence; It is the first Layer A concatenation matrix consisting of a set of first attention map matrices, It is The target matrix of the layer is used to convert The dimension is converted to The dimensions are the same, The initial network structure Layer A second attention map matrix; and is the target diagonal matrix. In these two matrices, the values ​​of elements other than the main diagonal elements are zero; It means taking the main diagonal elements of the matrix to form a column vector.

[0131] also,

[0132] in, Indicates the Layer, The distillation loss function of the group, is the objective function, which can be KL divergence; It is the first Layer A concatenation matrix consisting of a set of first attention map matrices, It is The target of the layer is used to The dimension is converted to The dimensions are the same, The initial network structure Layer A second attention map matrix; and is the target diagonal matrix. In these two matrices, the values ​​of elements other than the main diagonal elements are zero; It means taking the main diagonal elements of the matrix to form a column vector.

[0133] Based on the above overall distillation loss and task loss, the overall loss function is defined . The total distillation loss and mission loss is a function of the independent variable. In this embodiment, The linear function is as follows:

[0134] in, is a constant coefficient used to balance and Contributions, for example, By adjusting Better accuracy can be achieved.

[0135] In this embodiment, the overall distillation loss is constructed by calculating the difference between different network layers of the teacher network and the initial network structure. The differences between different network layers are taken into account, so that the initial network structure can learn richer feature representations and hierarchical knowledge structures from the teacher network. The overall distillation loss is combined with the task loss to construct an overall loss function to further guide the initial network structure to inherit the generalization ability of the teacher network. The initial network structure can better balance the inheritance of teacher network knowledge and the fitting of training data during the learning process, thereby improving the performance of the student network.

[0136] In some embodiments, the teacher network is a pre-trained model, and the student network is a lightweight target model constructed based on the pre-trained model; wherein the target model has the same number of network layers, the number of neurons in the network layer, and the activation function as the pre-trained model; and the number of attention heads in the target model is less than the number of attention heads in the pre-trained model.

[0137] In this embodiment, the pre-trained model is typically trained for a long time on a large dataset and possesses strong feature extraction and generalization capabilities. For example, in the field of natural language processing, large pre-trained models such as BERT (Bidirectional Encoder Representations from Transformers) and GPT (Generative Pre-trained Transformer) can be used as the teacher network. The structure of the teacher network is typically complex, containing many network layers, a large number of neurons, and a large number of attention heads.

[0138] The student network serves as a lightweight target model and is lightweighted based on the teacher network, retaining exactly the same number of network layers, the number of neurons in each layer, and the activation function. Parameter compression is achieved by reducing the number of attention heads (for example, from 12 to 4).

[0139] For example, consider a training dataset consisting of text data. This text data is fed into a pre-trained model, generating intermediate state values ​​during the model's execution, such as query embeddings, key embeddings, and value embeddings, as well as the final output text. In a pre-trained model for question answering, the input is the question, and the output is the answer.

[0140] The following example illustrates the training process of knowledge distillation. In this example, a large-scale text data set was collected, containing hundreds of millions of sentences, of which Chinese sentences accounted for 30%, English sentences accounted for 60%, and other languages ​​accounted for 10%. In the process of knowledge distillation, text data is sampled from the large-scale text data set by random sampling. The above language proportions are exemplary values, and the use of other numerical ratios will not affect the effect of the present invention.

[0141] To prevent the order of samples in the training dataset D from affecting the performance of the model, all samples in the training dataset D are randomly arranged at the beginning of knowledge distillation. In each round of training, the training algorithm reads a batch of samples, for example: 1024 samples.

[0142] Read data from the training dataset D, assuming that the data has been preprocessed as necessary and the format meets the model input requirements; initial training cycle , total training cycle ; The stochastic gradient descent algorithm is the Adam algorithm, and the relevant parameters are: the learning rate is , momentum coefficient , , The learning rate strategy is the cosine strategy, which is adjusted dynamically during the training cycle based on the cosine function. The number of samples in a mini-batch is B = 1024.

[0143] Target Model With pre-trained models The structure is basically the same, only the number of heads of the multi-head attention mechanism is different. The number of attention heads of the trained model is no greater than that of the pre-trained model.

[0144] The overall loss function for training is . Training steps: Step 1: Pre-trained model All parameter values ​​are fixed; according to the source model Build the target model based on the network structure The network structure of the target model is randomly initialized. The element values ​​of the weight matrix in , as well as other parameters; Step 2: When When , do the following: Step 2.1: ; Step 2.2: Train the dataset The order of samples in is randomly shuffled; Step 2.3: From the training dataset Select a batch B of training samples; Step 2.4: Based on the above training samples and parameter settings, calculate the loss value, and then use the stochastic gradient descent algorithm (Adam algorithm) to update the target model Before updating the parameters, the gradient is clipped and the gradient threshold is set to 5 (which can be adjusted according to the actual situation) to prevent gradient explosion.

[0145] Step 2.5: Repeat Step 2.3 and Step 2.4 until the training data set All samples were used.

[0146] Step 3: End training and output target model . Save the target model M' to the specified file path for subsequent use.

[0147] In this embodiment, by keeping the number of network layers, the number of neurons and the activation function unchanged, the student network can inherit the basic structure and information processing capabilities of the teacher network, and can learn the knowledge representation of the teacher network during the knowledge distillation process. The design of reducing the number of attention heads greatly reduces the computational complexity and resource consumption while maintaining the similarity of the model structure, making the model more suitable for deployment in resource-constrained environments.

[0148] The embodiment of the present application also provides an image retrieval system, including an image encoding module and a text encoding module; The image encoding module uses the image encoder in the pre-trained model as the teacher network and the student network obtained by any of the above knowledge distillation methods; The text encoding module is a student network obtained by using any of the above knowledge distillation methods, with the text encoder in the pre-trained model as the teacher network.

[0149] In the embodiment of the present application, the role of the image encoder is to convert the input image data into a feature vector with semantic meaning. The image encoder uses a deep learning network architecture to abstract and extract the pixel information in the image layer by layer, and can capture the visual features of the image such as shape, texture, and color. For example, when identifying a picture of a cat, the image encoder can extract features such as the cat's outline, the texture of the fur, and the shape of the eyes, and convert this complex visual information into image feature vectors that the computer can understand and process. These image feature vectors can subsequently be used for various visual tasks such as image classification, object detection, and image generation.

[0150] A text encoder processes text data and converts it into a vector representation containing semantic information. A text encoder can understand the meaning of words, phrases, sentences, and even paragraphs within a text. For example, for a text describing a landscape, a text encoder can identify key words such as "mountains," "rivers," and "beautiful," understand the relationships between these words, and generate a text feature vector that reflects the overall semantics of the text. This text feature vector can be used for natural language processing tasks such as text classification, sentiment analysis, and machine translation, enabling computers to understand and process text content in the same way humans do.

[0151] In some embodiments, the pre-trained model may be a CLIP (Contrastive Language-Image Pre-training) model. The CLIP model is a powerful multimodal pre-training model that includes an image encoder and a text encoder, and possesses excellent image and text encoding capabilities. In the CLIP model, the image encoder and text encoder are jointly trained to learn the semantic associations between images and text. For example, given a picture of a sunrise and a text describing "the sun rises from the horizon, radiant and brilliant," the CLIP model's image encoder and text encoder can extract features from the image and text, respectively, and determine a high degree of correlation between them through internal mechanisms.

[0152] In the embodiments of this application, Figure 4As shown in the figure, the image encoder in the CLIP model can be used as a teacher network to build a student network. The student network is trained by the above knowledge distillation method to obtain the image encoding module. The text encoder in the CLIP model is used as the teacher network to build a student network. The student network is trained using the above knowledge distillation method to obtain the text encoding module. .

[0153] Through this application, by grouping the attention graph matrix of the teacher network based on the number of attention heads in the student network, the high-dimensional attention space of the teacher network can be divided into subspaces of the same dimension as the student network. By splicing the attention graph matrices within the group, each group of spliced ​​matrices is aligned with the attention head of the student network, achieving a more regular correspondence between the attention graphs of the teacher network and the student network, making the attention graph matrices of the teacher and student networks comparable in form and structure, thereby being able to construct a distillation loss and accurately transfer the knowledge of the complex teacher network to the lightweight student network. Therefore, the difference in the number of heads of the teacher network and the student network in the multi-head attention mechanism can be solved, which hinders the student network from fully learning the knowledge of the teacher network, thereby achieving the technical effect of improving the knowledge distillation effect and enhancing the performance of the student network.

[0154] In some embodiments, the image encoding module is used to extract features from images in the image database and store the extracted image features in the image feature database; The text encoding module is used to extract features from the query text to obtain text features; calculate the similarity between the text features and each image feature in the image feature database, determine the target image features whose similarity with the text features is greater than a threshold, match the target image corresponding to the target image features from the image database, and output the target image.

[0155] In this embodiment, the image retrieval system can implement the following functionality: using a text query as a query, it retrieves one or more images from an image database that match the description in the text. For example, if the input query is "a cat," the image retrieval system will return images related to "cat," such as an orange cat napping in the sun or a black and white cat playing in the grass. Alternatively, if the input query is "beautiful scenery," the image retrieval system may return images of magnificent mountains and blue skies or stunning sunsets. Similarly, if the input query is "city nightscape," the image retrieval system may return images of shimmering city lights and tall buildings or busy streets at night. The image retrieval system helps users quickly find images that match their description through concise text queries, providing a convenient image search and browsing experience.

[0156] In this embodiment, the image database It is an image database to be queried, which includes a large number of unlabeled images. The content of these images is unknown, that is, the object categories in the images are unknown. The image database is large in scale and covers images from multiple fields and scenes. Each image in the image database is an independent instance and has not been classified or labeled in advance. For the convenience of subsequent description, the image database is defined as The total number of images is .

[0157] Image database Each image in Input to the image encoding module , get the output feature embedding vector for each image , .exist In, take As an image Features. Image feature database include image features, the set is , image features correspond one to one with images.

[0158] In this embodiment, the query text can be a natural language text, such as "Generate an image of a cat." It should be noted that the language type of the query text is consistent with the language type of the training dataset used by the training text encoding module. For example, if the text in the training dataset is in Chinese, the query text should also be in Chinese. If the text in the training dataset covers multiple languages, the input query text can include multiple languages.

[0159] Query text Input to the text encoding module The obtained feature embedding vector is , .exist In, take Serves as text features of the query text.

[0160] Image feature database include image features, the set is The text features of the query text T are , the similarity between the text feature of the query text T and each feature in the image feature data It can be calculated as:

[0161] in, ,in and Represent two column vectors, represent The transposed vector of Represents the vector modulus.

[0162] In this embodiment, the task of the image query system is to find one or more images corresponding to the query text content from the image database based on the input query text. To this end, a similarity threshold can be set ,For example, Of course, the similarity threshold Other values ​​are also possible and are not limited in the embodiments of the present application.

[0163] Based on the similarity threshold , and the image index set is , which means return all The set of indexes of , that is, the set of target image features.

[0164] Based on image index collection , from the image database Match the corresponding target image and return it.

[0165] The knowledge distillation method provided in the embodiment of the present application can be executed by a knowledge distillation device. In the embodiment of the present application, the knowledge distillation device is used as an example to illustrate the knowledge distillation method provided in the embodiment of the present application.

[0166] like Figure 5 As shown, the knowledge distillation device includes: A construction module 510 is used to construct an initial network structure of the student network based on the network structure of the teacher network; A grouping module 520 is used to group the multiple first attention map matrices of the teacher network; the number of groups is the same as the number of attention heads in the initial network structure; A splicing module 530 is used to splice the first attention map matrices in each group to obtain a splicing matrix corresponding to each group; The updating module 540 is used to calculate the concatenation matrix and the second attention map matrix of the initial network structure according to the distillation loss, and update the network parameters of the initial network structure based on the calculation results to obtain the student network.

[0167] Through this application, by grouping the attention graph matrix of the teacher network based on the number of attention heads in the initial network structure of the student network, the high-dimensional attention space of the teacher network can be divided into subspaces of the same dimension as the student network. By splicing the attention graph matrices within the group, each group of spliced ​​matrices is aligned with the attention heads of the initial network structure, achieving a more regular correspondence between the attention graphs of the teacher network and the student network, making the attention graph matrices of the teacher network and the student network comparable in form and structure, thereby being able to calculate the distillation loss and accurately transfer the knowledge of the complex teacher network to the lightweight student network. Therefore, the difference in the number of heads of the teacher network and the student network in the multi-head attention mechanism can be solved, which hinders the student network from fully learning the knowledge of the teacher network, thereby achieving the technical effect of improving the knowledge distillation effect and enhancing the performance of the student network.

[0168] In some embodiments, the grouping module 520 is further configured to: Calculate the ratio of the number of attention heads in the teacher network to the number of attention heads in the initial network structure; Determine a grouping strategy based on the ratio; Multiple first attention map matrices of the teacher network are grouped according to the grouping strategy.

[0169] In some embodiments, the grouping module 520 is further configured to: When the ratio is an integer, the grouping strategy includes: dividing the multiple first attention map matrices equally into group; among them, Equal to the number of attention heads in the student network.

[0170] In some embodiments, the grouping module 520 is further configured to: When the ratio is not an integer, multiple first attention map matrices are divided into group; of which, at most -1 The number of first attention maps in the group is the ratio rounded down; Equal to the number of attention heads in the student network.

[0171] In some embodiments, the update module 540 is further configured to: Reduce the dimension of the concatenated matrix so that the dimension of the concatenated matrix is ​​the same as the dimension of the second attention map matrix; The calculation results are obtained based on the concatenated matrix after dimensionality reduction and the second attention map matrix.

[0172] In some embodiments, the update module 540 is further configured to: The concatenated matrix is ​​reduced in dimensionality using a target matrix, where the target matrix is ​​a learnable matrix that is updated by a back-propagation algorithm during the initial network structure training process.

[0173] In some embodiments, the update module 540 is further configured to: The target diagonal matrix is ​​used for constraint, and the concatenated matrix after dimensionality reduction and the second attention map matrix are calculated based on the objective function to obtain the calculation result; the objective function represents the difference between the teacher network and the initial network structure.

[0174] In some embodiments, the update module 540 is further configured to: Get a labeled training dataset; The network parameters of the teacher network are fixed, the training dataset is input into the teacher network and the initial network structure, and the network parameters of the initial network structure are updated through the back propagation algorithm based on the overall loss function; the overall loss function is constructed based on the distillation loss and the task loss; the task loss represents the difference between the predicted value of the initial network structure and the true value corresponding to the label.

[0175] The description of the embodiment corresponding to the knowledge distillation device can be found in the relevant description of the embodiment corresponding to the knowledge distillation method, and will not be repeated here.

[0176] The embodiment of the present application also provides an electronic device, such as Figure 6 As shown, it includes a memory 602 and a processor 601, the memory 602 stores a computer program, and the processor 601 is configured to run the computer program to execute the steps in any of the above-mentioned knowledge distillation method embodiments.

[0177] An embodiment of the present application further provides a computer-readable storage medium, in which a computer program is stored, wherein the computer program is configured to execute the steps of any of the above-mentioned knowledge distillation method embodiments when run.

[0178] In an exemplary embodiment, the computer-readable storage medium may include, but is not limited to, various media that can store computer programs, such as a USB flash drive, a read-only memory (ROM), a random access memory (RAM), a mobile hard disk, a magnetic disk, or an optical disk.

[0179] An embodiment of the present application further provides a computer program product, which includes a computer program. When the computer program is executed by a processor, the steps in any of the above-mentioned knowledge distillation method embodiments are implemented.

[0180] An embodiment of the present application also provides another computer program product, including a non-volatile computer-readable storage medium, which stores a computer program. When the computer program is executed by a processor, it implements the steps in any of the above-mentioned knowledge distillation method embodiments.

[0181] Professionals may further appreciate that the units and algorithm steps of each example described in conjunction with the embodiments disclosed herein can be implemented in electronic hardware, computer software, or a combination of the two. In order to clearly illustrate the interchangeability of hardware and software, the above description has generally described the components and steps of each example according to their functions. Whether these functions are performed in hardware or software depends on the specific application and design constraints of the technical solution. Professionals and technicians may use different methods to implement the described functions for each specific application, but such implementation should not be considered beyond the scope of this application.

[0182] The above is a detailed introduction to a knowledge distillation method, electronic device and computer-readable storage medium provided by this application. This article uses specific examples to illustrate the principles and implementation methods of this application. The description of the above embodiments is only used to help understand the method of this application and its core idea. It should be pointed out that for ordinary technicians in this technical field, without departing from the principles of this application, several improvements and modifications can be made to this application, and these improvements and modifications also fall within the scope of protection of the claims of this application.

Claims

1. A knowledge distillation method, characterized in that: include: Construct the initial network structure of the student network according to the network structure of the teacher network; Grouping multiple first attention map matrices of the teacher network; the number of the groups is the same as the number of attention heads in the initial network structure; Splicing the first attention map matrices in each group to obtain a splicing matrix corresponding to each group; The concatenated matrix and the second attention map matrix of the initial network structure are calculated according to the distillation loss, and the network parameters of the initial network structure are updated based on the calculation results to obtain the student network.

2. The method according to claim 1, characterized in that The grouping of multiple first attention map matrices of the teacher network includes: Calculating the ratio of the number of attention heads in the teacher network to the number of attention heads in the initial network structure; determining a grouping strategy based on the ratio; The plurality of first attention map matrices of the teacher network are grouped according to the grouping strategy.

3. The method according to claim 2, characterized in that The grouping of the plurality of first attention map matrices of the teacher network according to the grouping strategy comprises: When the ratio is an integer, the grouping strategy includes: dividing the first attention map matrices into group; among them, is equal to the number of attention heads in the student network.

4. The method according to claim 2, characterized in that The grouping of multiple first attention map matrices of the teacher network includes: When the ratio is not an integer, the first attention map matrices are divided into group; of which, at most -1 The number of the first attention maps in the group is the ratio rounded down; is equal to the number of attention heads in the student network.

5. The method according to claim 1, wherein The calculating of the concatenation matrix and the second attention map matrix of the initial network structure according to the distillation loss includes: Performing dimensionality reduction on the splicing matrix so that the dimension of the splicing matrix after dimensionality reduction is the same as the dimension of the second attention map matrix; Calculation is performed based on the splicing matrix after dimensionality reduction and the second attention map matrix to obtain the calculation result.

6. The method according to claim 5, characterized in that The dimensionality reduction of the splicing matrix includes: A target matrix is ​​used to reduce the dimension of the concatenated matrix; wherein the target matrix is ​​a learnable matrix updated by a back-propagation algorithm during the initial network structure training process.

7. The method according to claim 5, characterized in that The calculation based on the splicing matrix after dimensionality reduction and the second attention map matrix includes: A target diagonal matrix is ​​used for constraint, and the concatenated matrix after dimensionality reduction and the second attention map matrix are calculated based on an objective function to obtain the calculation result; the objective function characterizes the difference between the teacher network and the initial network structure.

8. The method according to claim 7, characterized in that The target diagonal matrix includes a plurality of diagonal matrices that are updated by a back propagation algorithm during the initial network structure training process.

9. The method according to claim 7, characterized in that The distillation loss function is expressed as: in, Indicates the Layer, The distillation loss function of the group, represents the objective function, Indicates the The concatenation matrix of the group, represents the target matrix, represents the second attention map matrix, and represents different target diagonal matrices, represents the hyperparameter, It means taking the main diagonal elements of the matrix to form a column vector, Represents the zero norm of the matrix.

10. The method according to claim 1, characterized in that The step of calculating the concatenation matrix and the second attention map matrix of the initial network structure according to the distillation loss, and updating the network parameters of the initial network structure based on the calculation results to obtain the student network includes: Get a labeled training dataset; The network parameters of the teacher network are fixed, the training data set is input into the teacher network and the initial network structure, and the network parameters of the initial network structure are updated through a back propagation algorithm based on an overall loss function; wherein the overall loss function is constructed according to the distillation loss and the task loss; and the task loss represents the difference between the predicted value of the initial network structure and the true value corresponding to the label.

11. The method according to claim 10, characterized in that The task loss is expressed as: in, Indicates mission loss, Indicates the total number of words, represents the vocabulary size, yes The vector elements, The vector is dimension hot vector, indicating the true corresponding category at position i, yes The vector elements, Predict a V-dimensional probability distribution vector for the initial network structure, which indicates the probability that the word at position i predicted by the initial network structure belongs to each word in the vocabulary.

12. The method according to claim 10, characterized in that The overall loss function is constructed as follows: Constructing an overall distillation loss based on the distillation loss; the overall distillation loss represents the overall difference between different network layers of the teacher network and the initial network structure; Constructing the overall loss function according to the overall distillation loss and the task loss; The overall distillation loss is expressed as: in, represents the overall distillation loss, Representative Layer, The distillation loss of the group, Indicates the number of attention heads in the initial network structure; The overall loss function is expressed as: in, represents the overall loss function, Indicates mission loss, represents the overall distillation loss, is a constant coefficient.

13. The method according to claim 1, wherein The teacher network is a pre-trained model, and the student network is a lightweight target model constructed based on the pre-trained model; wherein the target model has the same number of network layers, the number of neurons in the network layer, and the activation function as the pre-trained model; and the number of attention heads in the target model is less than the number of attention heads in the pre-trained model.

14. An electronic device, characterized in that: include: memory for storing computer programs; A processor, configured to implement the steps of knowledge distillation as claimed in any one of claims 1 to 13 when executing the computer program.

15. A computer-readable storage medium, characterized in that The computer-readable storage medium stores a computer program, wherein the computer program, when executed by a processor, implements the steps of the knowledge distillation method according to any one of claims 1 to 13.

Images