A modular decomposition method for transformer neural network models

Abstract

The application provides a modular decomposition method for a Transformer neural network model, modular training is performed on a randomly initialized model, and structured decomposition is implemented by applying a mask after the modular training, so that flexible on-demand reuse is realized; specifically comprising the following steps: S1, modular training; given a neural network model, firstly, randomly initialize all parameters of the model; a neuron recognizer is included, which is used to identify neurons related to a specific function; S2, structured decomposition; the modular training model is subjected to modular decomposition; S3, on-demand reuse; on-demand reuse is realized by structurally removing neurons, and memory and computing overheads are minimized. The application can realize structured decomposition of a Transformer structure model and flexible on-demand reuse of functions, and has high expandability.

Description

Technical Field

[0001] This invention relates to the field of neural network model reuse technology in computer intelligent software engineering, specifically a modular decomposition method for Transformer neural network models. Background Technology

[0002] In recent years, deep learning technology has flourished, and neural network models have demonstrated unique advantages in various fields such as image classification, object detection, speech recognition, and natural language processing. However, as neural network models become increasingly complex and large, the overhead of model building (such as training time and computational resource costs) has also increased dramatically, placing a heavy burden on model developers. Fortunately, a vast number of pre-trained neural network models have been publicly released in various open-source communities. For example, the NLP open-source community Hugging Face has more than 80,000 publicly available pre-trained models, and there are as many as 90,000 repositories tagged with machine learning on GitHub. Inspired by software reuse, reusing existing pre-trained neural network models to reduce (or even eliminate) the high cost of model building has been widely researched and applied.

[0003] To achieve the reuse of pre-trained models, the industry generally adopts two types of techniques: model selection and network-based deep transfer learning.

[0004] Model selection techniques can quickly identify pre-trained models that can be directly reused, while network-based deep transfer learning techniques can train new models for the target task with minimal time cost and labeled data. However, both of these techniques focus more on whether the performance of the reused model meets the requirements of the target task, neglecting the reuse overhead, including the training and inference overhead of the model on the target task data.

[0005] Therefore, in order to solve these problems, there are still immature technologies in the industry:

[0006] Prior Art 1, for the N-class CNN model (Binhang Qi, Hailong Sun, Hongyu Zhang, and Xiang Gao. 2024. Reusing Convolutional Neural Network Models through Modularization and Composition. ACM Trans. Softw. Eng. Methodol. 33, 3, Article 70 (Mar 2024), 39 pages.), this technology proposes a compressed convolutional neural network modularization algorithm and applies it to the quality assurance of CNN models. For each of the N classes, the algorithm forms a module by searching for the convolutional kernels in the pre-trained model that are used to extract the features of that class and removing other irrelevant convolutional kernels.

[0007] Prior Art 1 can effectively achieve the modularization of CNN models. At the same time, for a single class, since the generated module only retains the partial convolutional kernels related to that class in the pre-trained model, the number of parameters of the module is smaller than that of the pre-trained model. Compared with other modularization technologies, the combined model composed of the modules generated by this technology has a smaller number of parameters and lower inference overhead. However, this technology still has the following deficiencies: First, the modularization method of this technology belongs to post-training modularization. It decomposes the already trained model. However, since the original model does not consider the modularization feature during the training process, the different convolutional kernels in the model are highly coupled, which greatly limits the modularization effect; Second, this technology targets the convolutional kernel as a sub-structure for modularization, which makes this technology only applicable to CNN and cannot be extended to other models.

[0008] Prior Art 2, CN116340808A discloses a modularization method for neural network models. For an input multi-class pre-trained neural network model and an image or text multi-class dataset for a target task, it is processed by a processing unit, which includes three parts: a search space, a performance evaluation strategy, and a search strategy; the specific inputs of the method include: a pre-trained image or text multi-class model as the original model, which is composed of several neural network layers and is an N-class model with a total of L weights; and an image or text K-class dataset for a target task, where K < N, including data for modularization and data to be classified; the method realizes the modularization of the original model through six steps, generates a K-class neural network module for the target task that only contains J weights, where J < L, and further uses the neural network module to classify the data to be classified for the target task, and finally outputs the classified image or text data.

[0009] Existing technique two effectively modularizes CNN models and can be extended to other models. Furthermore, for a single category, since the generated module only retains the weights relevant to that category from the pre-trained model, the number of parameters in the module is smaller than in the pre-trained model. Compared to other modularization techniques, this technique has better scalability because it sets the target as "weights," a unit universally present in all models. However, this technique still has the following shortcomings: First, its modularization method is still post-training modularization, decomposing the already trained model. However, since the original model did not consider modularity during training, the different convolutional kernels are highly coupled, greatly limiting the modularization effect. Second, although this technique targets weights and achieves the goal of setting a large number of weights to 0 to obtain a sparse weight matrix, it cannot structurally remove these weights. Currently, the inference acceleration of the model on the sparse weight matrix relies on a specific sparse acceleration engine, which is difficult to accelerate on high-parallel devices such as GPUs, and therefore cannot be applied in real-world production environments.

[0010] The existing technical solution (Binhang Qi, Hailong Sun, Hongyu Zhang, Ruobing Zhao, and Xiang Gao. 2024. Modularizing while Training: A New Paradigm for Modularizing DNN Models. In Proceedings of the IEEE / ACM 46th International Conference on Software Engineering (ICSE'24). Association for Computing Machinery, New York, NY, USA, Article 31, 12 pages.) proposes a modularization algorithm for convolutional neural networks during training for N-class CNN models, and applies it to the modular decomposition and on-demand reuse of CNN models. During training, for each of the N classes, the algorithm adds a kernel recognizer to each convolutional neural network layer to identify which kernels are related to the target class, and defines cohesion and coupling losses. Finally, the model is trained to have high cohesion and low coupling characteristics, and modules are obtained by structurally removing convolutional kernels that are irrelevant to the corresponding class.

[0011] Existing technique three can effectively achieve modularization of CNN models during training. Furthermore, for a single category, since the generated modules only retain the convolutional kernels relevant to that category from the pre-trained model, the number of parameters in each module is smaller than that of the pre-trained model. Compared to other modularization techniques, because this technique introduces modularity constraints during the model training phase, the combined model composed of modules generated by this technique has fewer parameters and lower inference overhead. However, this technique still has the following shortcomings: First, this technique targets the convolutional kernel as a substructure for modularization, which makes it only applicable to CNNs and cannot be generalized to other models. Second, the optimization function used in this technique is the weighted sum of cohesion loss, coupling loss, and cross-entropy loss, which presents some difficulties in multi-objective optimization, specifically manifested in decreased accuracy and room for further reduction in the convolutional kernel retention rate. Summary of the Invention

[0012] This invention addresses the problem of on-demand reuse of deep neural network models, aiming to provide a modular decomposition method for neural network models with Transformer structures, supporting modular decomposition during model training at the neuron granularity.

[0013] Specifically, the main problems to be solved include:

[0014] (1) Existing modularization methods during training can only be implemented for CNN models and cannot be applied to the current state-of-the-art Transformer model structure. This invention achieves modularization on the Transformer model by defining the modularization granularity at the neuron level and designing a neuron recognizer to determine the function corresponding to the neuron.

[0015] (2) A structured decomposition based on neurons was designed and implemented. Considering that a neuron is an abstract concept in the DNN model, which represents the output of a row of the weight matrix of the corresponding layer, a method for structured removal of the weight matrix in the calculation of attention mechanism and fully connected layer was designed.

[0016] (3) A loss function suitable for larger-scale model optimization was designed, which integrates cohesion, coupling and contrastive learning methods, and has made significant progress in the neuron retention rate, accuracy, cohesion and coupling of modular training.

[0017] The specific technical solution is as follows:

[0018] A modular decomposition method for Transformer neural network models is proposed, which involves modular training of a randomly initialized model and applying a mask after modular training to achieve structured decomposition, thereby enabling flexible on-demand reuse.

[0019] S1, Modular Training.

[0020] Given a neural network model, first initialize all parameters of the model randomly. Include a neuron recognizer to identify neurons associated with specific functions.

[0021] Specifically, the neuron recognizer comprises several mask generators, each consisting of multiple linear layers, used to learn the relationship between input features and corresponding categories. Each mask generator corresponds to a linear layer in the neural network model, and its input is identical to that of the corresponding linear layer. During the forward propagation of the model, the mask generator produces a mask with the same dimension as the number of neurons in the corresponding linear layer, where each value represents whether the corresponding neuron is associated with a category. The mask is used to calculate cohesion and coupling losses, and is ultimately optimized in the contrastive learning loss. After the forward propagation of the model, the cross-entropy loss is calculated on the final classification result. The objective function is a weighted sum of the cross-entropy loss, cohesion loss, and coupling loss.

[0022] S2, Structured Decomposition.

[0023] For models trained in a modular fashion, they will be decomposed into modular components.

[0024] Specifically, based on the mask generated by the mask generator, it is determined which neurons are relevant to the current target classification. Considering that a neuron actually corresponds to the output of a row in the weight matrix of that linear layer, the corresponding row in the weight matrix is ​​removed based on the mask. Furthermore, considering that the weight matrix of the previous layer is also partially removed, the corresponding row or column of the current layer is removed based on the mask of the previous layer. After decomposition, a decomposed module is obtained, and a target task classification head is added for simple fine-tuning, resulting in the final module for on-demand reuse.

[0025] S3, Reuse on Demand

[0026] By structurally removing neurons, on-demand reuse can be achieved, minimizing memory and computational overhead.

[0027] Specifically:

[0028] S1, Modular Training

[0029] Training a modular DNN model involves three steps: (1) identifying relevant neurons, (2) evaluating modularity and accuracy metrics, and (3) optimizing modularity and classification accuracy performance.

[0030] (1) Identifying Relevant Neurons. A neuron recognizer is used to identify neurons related to the target category in the linear and attention layers of the Transformer structure. For a model, its neuron recognizer I... DNN It contains several mask generators G, each of which is a fully connected neural network consisting of two linear layers, and attached to a linear layer L of the model. DNN Generate a mask m and apply m to the linear layer L. DNN The output h is then used to obtain the output after it has been masked. The process is represented by Formula 1 as follows:

[0031]

[0032] Equation 2 performs average pooling on the input x to generate a downsampled feature x. pooled And it is used as input to the mask generator G. In Equation 3, G encodes x pooled And generate a mask m, representing the value in L. DNN Which neuron corresponds to the current input sample?

[0033] In the ViT model, the feature dimensions of the input x are (B, N, D), where B represents the batch size of the input data, N represents the number of patches in the ViT model, and D represents the feature size of each patch. During forward propagation, training data is input into the model and the neuron recognizer. DNN Output a feature h with dimensions (B, N, D), and the neuron recognizer is for each L. DNN A mask m is generated, whose dimensions (B, 1, D) correspond to the output of each neuron. In Equation 4, for the output of a neuron, the mask m is broadcast along the dimensions of the patch and multiplied to obtain the masked output. G maps the value of m to the interval [0, 1) by adding the activation functions Tanh and ReLU. 0 represents that the corresponding neuron is independent of the input sample.

[0034] (2) Evaluation of modularity and classification performance. The concepts of cohesion and coupling are introduced to evaluate the performance of modularity. Combining the concepts of cohesion and coupling, positive samples are constructed as the overlap of activated neurons of data samples belonging to the same class, i.e., cohesion. Similarly, negative samples are constructed as the overlap of activated neurons of data samples belonging to different classes, i.e., coupling.

[0035] Specifically, in a dataset containing n categories, for category c i It belongs to category c i n i Each sample is represented as Correspondingly, the set of neurons activated for each sample is represented as For category c i Module mN i Includes neurons used by all samples For category c i Module mN i The degree of cohesion is expressed as:

[0036]

[0037] Coupling degree is expressed as:

[0038] Coupling(m i m j =Overlap(mN) i mN j 6

[0039] The overlap metric is derived from JaccardIndex. For two sets A and B, their JaccardIndex is:

[0040]

[0041] (3) Modular training and optimization. Cosine similarity is used instead of overlap to obtain an optimized cohesion and coupling loss:

[0042]

[0043] n i This indicates that within each batch, those belonging to category c i The number of samples. This indicates that it belongs to category c. i The mask for the j-th sample. Cos(a, b) represents the cosine similarity between two tensors a and b. A temperature parameter τ and an exponential function from contrastive learning are introduced to scale the distribution of samples in the feature space; the cohesion loss is expressed as:

[0044]

[0045] The loss of coupling is expressed as:

[0046]

[0047] After obtaining the cohesion and coupling losses, the contrast loss is defined as follows;

[0048]

[0049] The target loss function consists of two parts: cross-entropy loss and contrastive loss.

[0050]

[0051] Here, α represents the weight of the contrastive loss. Based on the above loss function, the model is optimized using the gradient descent algorithm to train the corresponding neurons to recognize samples of a specific class.

[0052] S2, Structured Decomposition

[0053] The modular model is decomposed into multiple modules based on the mask generated by the mask generator. Specifically, it is first determined which part of the model a neuron corresponds to, and then the model is decomposed based on the mask.

[0054] (1) Determine the neuron structure.

[0055] The goal of removing neurons is achieved by removing a specific row from the weight matrix, thereby reducing the corresponding computational overhead.

[0056] (2) Decompose the linear layer using a partial weight matrix.

[0057] To decompose the modular model into modules, a mask is first generated for each sample. For class c i n i Each belongs to c i The samples are represented as Their corresponding masks are represented as definition:

[0058]

[0059] For category c i In order to select neurons that are effective for all relevant samples and to filter out noise, the module mask...

[0060]

[0061] For a fully connected layer, for two consecutive layers A and B, B receives the output of A. Therefore, for the weight matrix B, the weight rows where the mask B is 0 need to be removed in the output dimension. Simultaneously, since B receives the output of A, the weight matrix B also needs to remove the weight columns where the mask A is 0 in the input dimension. For the attention layer, there are three weight matrices: Query, Key, and Value. All three take the same features as input and obtain the corresponding Q, K, and V. Then, the attention score is calculated as follows:

[0062]

[0063] Where Q = [q1, q2, ..., q n ], K = [k1, k2, ..., k n ], For Query and Key, take the intersection of their masks, i.e., m QK =m Q ∩m K Only the common parts are retained. QK To remove.

[0064] When the linear layers at both ends of the residual connection are removed in a structured manner, the output and input dimensions change. An on-demand fill-in and patching solution is adopted to fill the input brought by the residual connection according to the input dimension of the subsequent layer of the residual connection and remove the redundant part.

[0065] The beneficial effects of the technical solution of this invention are as follows:

[0066] 1. This invention provides a modular decomposition method for neural network models with Transformer structures, enabling structured decomposition of Transformer models and flexible, on-demand reuse of functions. Compared to the unstructured method of technique two, neurons in the model can be physically removed using the method of this invention, thereby reducing storage and inference overhead.

[0067] 2. This invention selects neurons as the modular granularity and provides a complete method for the structural decomposition of neurons. Compared to techniques one and three, this granularity is widely present in all models and has high scalability.

[0068] 3. The neuron structure decomposition method designed in this invention retains as few as 18.8% of the neurons in the original model (corresponding to 12.04% weight) in the module obtained on the ViT model. At the same time, its classification accuracy is comparable to other methods, with almost no decrease.

[0069] 4. The loss function designed in this invention significantly reduces the retention rate of neurons or convolutional kernels. Ablation tests on CNN models (using the same CNN model as in Technique 3, but not the ViT model of this invention, only the loss function was replaced) show that the average convolutional kernel retention rate of the method of this invention is 12.8%, while the average convolutional kernel retention rate of the method in Technique 3 is 17.55%. Attached Figure Description

[0070] Figure 1 This is a flowchart of the present invention;

[0071] Figure 2 This is a schematic diagram of the removal of neurons in the attention and fully connected layers of the present invention;

[0072] Figure 3 This is a schematic diagram of the dimension mismatch and repair of residual connections in this invention;

[0073] Figure 4This is a schematic diagram illustrating the on-demand reuse of modules in this invention. Detailed Implementation

[0074] This invention designs and implements NeMo, a modular method for structured training at the neuron granularity. The core workflow framework of this invention is as follows: Figure 1 As shown, NeMo mainly includes two processes: modular training of a randomly initialized model, and applying a mask after modular training to achieve structured decomposition, thereby enabling flexible on-demand reuse.

[0075] S1, Modular Training.

[0076] Given a neural network model, NeMo first randomly initializes all the model's parameters. NeMo includes a neuron recognizer to identify neurons associated with specific functions. Specifically, NeMo's neuron recognizer contains several mask generators, each consisting of multiple linear layers, used to learn the relationship between input features and corresponding categories. Each mask generator corresponds to a linear layer in the neural network model, and its input is identical to the input of the corresponding linear layer. During the model's forward propagation, the mask generator produces a mask with the same dimension as the number of neurons in the corresponding linear layer, where each value represents whether the corresponding neuron is associated with a category (0 indicates no association, 1 indicates association). The mask is used to calculate cohesion and coupling losses, which are ultimately optimized in the contrastive learning loss. After the model's forward propagation, the cross-entropy loss is calculated on the final classification result. The objective function is a weighted sum of the cross-entropy loss, cohesion loss, and coupling loss. As training progresses and backpropagation optimizes the losses, the model is ultimately trained into a modular model with high accuracy, high cohesion, and low coupling.

[0077] S2, Structured Decomposition.

[0078] For modularly trained models, NeMo performs modular decomposition. Specifically, NeMo determines which neurons are relevant to the current target classification based on the mask generated by the mask generator. Considering that a neuron actually corresponds to the output of a row in the weight matrix of that linear layer, NeMo removes that row from the weight matrix based on the mask. Furthermore, considering that the weight matrix of the previous layer is also partially removed, NeMo also removes the corresponding row or column of the current layer based on the mask of the previous layer, depending on the connection method between different layers, such as attention layers or fully connected layers. After decomposition, NeMo obtains a decomposed module, adds a target task classification head for simple fine-tuning, and obtains the final module for on-demand reuse.

[0079] The following details the steps described above:

[0080] S1, Modular Training

[0081] Training a modular DNN model mainly involves three steps: (1) identifying relevant neurons, (2) evaluating modularity and accuracy metrics, and (3) optimizing modularity and classification accuracy performance. Finally, the modularly trained model can be decomposed into smaller modules and reused flexibly as needed.

[0082] (1) Identifying Relevant Neurons. Identifying neurons associated with a specific category is a crucial part of modular training. This invention proposes a neuron recognizer in NeMo for identifying neurons associated with the target category in linear and attention layers of a Transformer architecture. For a model, its neuron recognizer I... DNN It contains several mask generators G, each of which is a fully connected neural network consisting of two linear layers, and attached to a linear layer L of the model. DNN Generate a mask m and apply m to the linear layer L. DNN The output h is then used to obtain the output after it has been masked. The process is represented by Formula 1 as follows:

[0083] To reduce additional training overhead, this invention performs average pooling on the input x in Equation 2, generating a downsampled feature x. pooled And it is used as input to the mask generator G. In Equation 3, G encodes x pooled And generate a mask m, representing the value in L. DNN Which neuron corresponds to the current input sample? Specifically, in the ViT model, the feature dimensions of the input x are (B, N, D), where B represents the batch size of the input data, N represents the number of patches in the ViT model, and D represents the feature size of each patch. For example, for an input image of size 224×224 divided into 14×14 patches, each patch is 16×16 in size. Therefore, N = 196 + 1 (an extra patch is used to represent global features), and D = 256. During forward propagation, training data is input into the model and the neuron recognizer. DNN It will output a feature h with dimensions (B, N, D). The neuron recognizer will assign a feature h to each L. DNNA mask m is generated, whose dimensions (B, 1, D) correspond to the output of each neuron. In Equation 4, for the output of a neuron, the mask m is broadcast along the dimensions of the patch and multiplied to obtain the masked output. G maps the value of m to the interval [0, 1) by adding the activation functions Tanh and ReLU. 0 represents that the corresponding neuron is independent of the input sample.

[0084] (2) Evaluation of modularity and classification performance. In modular training, this invention introduces the concepts of cohesion and coupling to evaluate the performance of modularity. Cohesion measures the degree of overlap of neurons used by samples with similar labels; coupling measures the degree of overlap of neurons used by samples with different labels. To optimize these two metrics, the contrastive learning method focuses on distinguishing between positive and negative samples, aiming to make positive samples closer in the feature space and negative samples farther apart. Existing methods for constructing positive samples mainly involve data augmentation of a single image to obtain a series of positive samples. Combining the concepts of cohesion and coupling, this invention constructs positive samples as the overlap of activation neurons of data samples belonging to the same class, i.e., cohesion. Similarly, negative samples are the overlap of activation neurons of data samples belonging to different classes, i.e., coupling. Specifically, in a dataset containing n classes, for class c i It belongs to category c i n i Each sample is represented as Correspondingly, the set of neurons activated for each sample can be represented as For category c i Module mN i Includes neurons used by all samples For category c i Module mN i The degree of cohesion can be expressed as:

[0085]

[0086] Coupling degree can be expressed as:

[0087] Coupling(m i m j =Overlap(mN) i mN j ) twenty one

[0088] The overlap metric is derived from JaccardIndex, which is widely used to calculate the overlap between two modules. For two sets A and B, their JaccardIndex is:

[0089]

[0090] (3) Modular Training and Optimization. Since the calculation of cohesion and coupling is non-differentiable, in order to continuously optimize modular performance during training, this invention uses cosine similarity instead of overlap, thereby obtaining an optimizable cohesion and coupling loss.

[0091]

[0092] Here, n i This indicates that within each batch, those belonging to category c i The number of samples. This indicates that it belongs to category c. i The mask for the j-th sample. Cos(a, b) represents the cosine similarity between two tensors a and b. This invention introduces the temperature parameter τ and an exponential function from contrastive learning to scale the distribution of samples in the feature space, thereby better focusing on samples with difficult-to-optimize margins and achieving lower retention rate and higher accuracy. Finally, the cohesion loss is expressed as:

[0093]

[0094] The loss of coupling is expressed as:

[0095]

[0096] After obtaining the cohesion and coupling losses, the present invention defines the contrast loss as follows;

[0097]

[0098] Furthermore, to optimize the classification accuracy of the model, this invention uses cross-entropy loss. The target loss function consists of two parts: cross-entropy loss and contrastive loss.

[0099]

[0100] Here, α represents the weight of the contrastive loss. Based on the above loss function, the model can be optimized using the gradient descent algorithm to train the corresponding neurons to recognize samples of a specific class.

[0101] S2, Structured Decomposition

[0102] To obtain the module corresponding to a category sample, after modular training, this invention decomposes the modular model into multiple modules based on the mask generated by the mask generator. Specifically, this invention first determines which part of the model a neuron corresponds to, and then decomposes the model based on the mask.

[0103] (1) Determining the Neuron Structure. Neurons are widely present in deep neural networks (DNNs). However, they are more of a descriptive concept than a concrete structure. In a linear layer, a neuron represents the final value after the input has passed through a weight matrix, a bias has been added, and an activation function has been applied. Therefore, unlike a convolutional kernel, a neuron is not a concrete structure. However, this invention can remove neurons by removing a corresponding row from the weight matrix, thereby reducing the corresponding computational overhead.

[0104] (2) Decomposing the linear layer using a partial weight matrix. To decompose the modular model into modules, NeMo first generates a mask for each sample. For class c i n i Each belongs to c i The sample can be represented as Their corresponding masks can be represented as This invention defines:

[0105]

[0106] For category c i In order to select neurons that are effective for all relevant samples and to filter out noise, the module mask...

[0107]

[0108] For fully connected layers, weight removal is straightforward. For two consecutive layers A and B, where B receives the output of A, the weight matrix B needs to have its weight rows (with B as 0) removed from the output dimension. Simultaneously, since B receives the output of A, the weight matrix B also needs to have its weight columns (with A as 0) removed from the input dimension, as follows: Figure 2 The blue section shows the situation. For the attention layer, the situation is much more complex. The attention mechanism calculation involves three weight matrices: Query, Key, and Value. All three receive the same features as input and obtain the corresponding Q, K, and V. Then, the attention score is calculated as follows:

[0109]

[0110] Where Q = [q1, q2, ..., q n ], K = [k1, k2, ..., k n ], When q i After being masked by 0, the corresponding q i ×k i A value of 0 represents k i The calculation is also meaningless. Therefore, for Query and Key, the intersection of their masks should be taken, i.e., mQK =m Q ∩m K Only the common parts are retained. For Value, the situation is similar to that of a fully connected layer. The output dimension of the weight matrix V can be determined by m... v To remove, but since it accepts the result of multiplying Q and K, its input dimension should be determined by m. QK To remove, see Figure 2 The yellow part.

[0111] Furthermore, since residual connections are widespread in neural networks, their abrupt connections between two different blocks introduce additional problems for modularity. Specifically, when the linear layers at both ends of a residual connection are structurally removed, the output and input dimensions change. Although the previous section of this invention addressed consecutive layers to ensure correct dimensionality matching after removal, residual connections typically span multiple layers or blocks, resulting in dimensionality mismatch at their ends (see...). Figure 3 As shown in (b) of the diagram.

[0112] Existing solutions include Figure 3 (c) Adding extra neurons for repair, and Figure 3 In the diagram (d), the outputs of the preceding and following layers are padded with zeros. However, these methods all incur additional overhead. Considering that when a portion of the weights of the following layer are removed, its input only needs to consider the retained portion, this invention proposes... Figure 3 The solution in (e) is to fill in the input as needed based only on the input dimension of the residual connection, removing the redundant parts. Clearly, this is the most efficient solution.

[0113] S3, Reuse on Demand

[0114] NeMo achieves efficient and flexible on-demand reuse by structurally removing neurons, minimizing memory and computational overhead. For example, when a task requires identifying "cat" or "dog," NeMo can decompose a 2-classification module from a 10-classification model and apply it, reducing inference overhead by removing extra irrelevant neurons.

[0115] like Figure 4 As shown, after modular training, developers can share the mask file to a model sharing platform. Users can obtain the corresponding module by applying the mask and obtain a high-accuracy module through simple fine-tuning.

[0116] This invention provides a modular training method at the neuron level and offers a detailed solution for the ViT model. For other structurally similar models such as DeiT, this method can also be used to implement modular training and structured decomposition processes, achieving excellent results.

[0117] The "neuron recognizer" used in this invention includes multiple "mask generators," each consisting of multiple linear layers. In the technical implementation, the linear layers within the mask generator can be replaced with convolutional layers or other network layers, all with the purpose of outputting a mask to identify activated neurons.

Claims

1. A modular decomposition method for Transformer neural network models, characterized in that, Modular training is performed on a randomly initialized model, and a mask is applied after modular training to achieve structured decomposition, thereby enabling flexible on-demand reuse. Specifically, the following steps are included: S1, Modular Training; Given a neural network model, the ViT model, and an image as input, first randomly initialize all parameters of the model; include a neuron recognizer to identify neurons associated with specific functions; S2, Structured decomposition; For models trained modularly, they will be decomposed into modules. The specific method of structured decomposition is as follows: based on the mask generated by the mask generator, determine which neurons are related to the current target classification; considering that a neuron actually corresponds to the output of a certain row of the weight matrix in the linear layer, remove the corresponding row of the weight matrix based on the mask; in addition, considering that the weight matrix of the previous layer will also be partially removed, remove the corresponding row or column of the current layer based on the mask of the previous layer; after decomposition, a decomposed module is obtained, and a target task classification head is added for simple fine-tuning to obtain the final module for on-demand reuse; The modular model is decomposed into multiple modules based on the mask generated by the mask generator; specifically, it is first determined which part of the model a neuron corresponds to, and then the model is decomposed based on the mask. (1) Determine the neuron structure; The goal of removing neurons is achieved by removing a specific row from the weight matrix, thereby reducing the corresponding computational overhead. (2) Decompose the linear layer using a partial weight matrix; To decompose the modular model into modules, a mask is first generated for each sample; for class c i n i Each belongs to c i The samples are represented as Their corresponding masks are represented as definition: For category c i In order to select neurons that are effective for all relevant samples and to filter out noise, the module mask... For a fully connected layer, for two consecutive layers A and B, B receives the output of A; therefore, for the weight matrix B, it needs to remove the weight rows where the mask B is 0 in the output dimension; simultaneously, since B receives the output of A, the weight matrix B also needs to remove the weight columns where the mask A is 0 in the input dimension; for the attention layer, it contains three weight matrices Query, Key, and Value, all of which take the same features as input and obtain the corresponding Q, K, and V; then, the attention score is calculated as follows: Where Q = [q1, q2, ..., q n ], K = [k1, k2, ..., k n ], For Query and Key, take the intersection of their masks, i.e., m QK =m Q ∩m K Only the common parts are retained; the output dimension of the weight matrix V is determined by m. QK To remove; When the linear layers at both ends of the residual connection are removed in a structured manner, the output and input dimensions change. An on-demand fill-in and patching solution is adopted to fill in the input brought by the residual connection according to the input dimension of the subsequent layer of the residual connection and remove the redundant part. S3. Reuse on demand; By structurally removing neurons, on-demand reuse can be achieved, minimizing memory and computational overhead.

2. The modular decomposition method for Transformer neural network models according to claim 1, characterized in that, S1. Specific methods of modular training: The neuron recognizer comprises several mask generators, each consisting of multiple linear layers, used to learn the relationship between input features and corresponding categories. Each mask generator corresponds to a linear layer in the neural network model, and its input is identical to that of the corresponding linear layer. During the forward propagation of the model, the mask generator produces a mask with the same dimension as the number of neurons in the corresponding linear layer, where each value represents whether the corresponding neuron is associated with a category. The mask is used to calculate cohesion and coupling losses, and is ultimately optimized in the contrastive learning loss. After the forward propagation of the model is completed, the cross-entropy loss is calculated on the final classification result. The objective function is a weighted sum of cross-entropy loss, cohesion loss, and coupling loss.

3. The modular decomposition method for Transformer neural network models according to claim 2, characterized in that, Training a modular DNN model in S1 involves three steps: (1) identifying relevant neurons, (2) evaluating modularity and accuracy metrics, and (3) optimizing modularity and classification accuracy performance. (1) Identify relevant neurons; a neuron recognizer is used to identify neurons related to the target category in the linear layer and attention layer of the Transformer structure; for a model, its neuron recognizer L DNN It contains several mask generators G, each of which is a fully connected neural network consisting of two linear layers, and attached to a linear layer L of the model. DNN ; Generate mask m and apply m to linear layer Lx NN The output h is then used to obtain the output after it has been masked. The process is represented by Formula 1 as follows: Equation 2 performs average pooling on the input x to generate a downsampled feature x. pooled And used as input to the mask generator G; In Formula 3, G encodes x pooled And generate a mask m, representing the value in L. DNN Which neuron corresponds to the current input sample? In the ViT model, the feature dimensions of the input x are (B, N, D), where B represents the batch size of the input data, N represents the number of patches in the ViT model, and D represents the feature size of each patch; during forward propagation, training data is input into the model and the neuron recognizer; L DNN Output a feature h with dimensions (B, N, D), and the neuron recognizer is for each L. DNN A mask m is generated, whose dimension (B, 1, D) corresponds to the output of each neuron; in Equation 4, for the output of a neuron, the mask m is broadcast along the dimension of the patch and multiplied to obtain the masked output; G maps the value of m to the interval [0, 1) by adding activation functions Tanh and ReLU; 0 represents that the corresponding neuron is independent of the input sample; (2) Evaluation of modularity and classification performance; The concepts of cohesion and coupling are introduced, and the performance of modularity is evaluated accordingly; Combining the concepts of cohesion and coupling, positive samples are constructed as the overlap of activation neurons of data samples belonging to the same class, i.e., cohesion; Similarly, negative samples are the overlap of activation neurons of data samples belonging to different classes, i.e., coupling. Specifically, in a dataset containing n categories, for category c i It belongs to category c i n i Each sample is represented as Correspondingly, the set of neurons activated for each sample is represented as For category c i Module mN i Includes neurons used by all samples For category c i Module mN i The degree of cohesion is expressed as: Coupling degree is expressed as: Coupling(m i ,m j )=Overlap(mN i ,mN j ) 6 The overlap index is derived from JaccardIndex; for two sets A and B, their JaccardIndex is: (3) Modular training and optimization; using cosine similarity instead of overlap to obtain an optimized cohesion and coupling loss: n i This indicates that within each batch, those belonging to category c i The number of samples; This indicates that it belongs to category c. i The mask for the j-th sample; Cos(a,b) represents the cosine similarity between two tensors a and b; a temperature parameter τ and an exponential function are introduced from contrastive learning to scale the distribution of samples in the feature space; the cohesion loss is expressed as: The loss of coupling is expressed as: After obtaining the cohesion and coupling losses, the contrast loss is defined as follows; The target loss function consists of two parts: cross-entropy loss and contrastive loss. Here, α is the weight of the contrastive loss; based on the above loss function, the model is optimized by the gradient descent algorithm, thereby training the corresponding neurons to recognize specific class samples.

Images