Distributed model training method based on neural network key layer
By determining the key layers in distributed model training and performing gradient sparse transmission and error accumulation, the data transmission bottleneck and model matching problems are solved, training efficiency and accuracy are improved, and low-cost and efficient distributed training is achieved.
Patent Information
- Application Number
- CN202510580174.8
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2025-05-07
- Publication Date
- 2025-08-01
AI Technical Summary
The existing distributed model training methods have problems such as data transmission bottlenecks and mismatch of the model with the data set, resulting in inefficient training, and existing solutions such as neural network architectures are expensive to search and lack practical feasibility.
The distributed model training method based on the key layer of neural network is adopted. By determining the key layer on a single machine, gradient sparse transmission and error accumulation are performed in multi-machine distributed deployment, combined with the search technology of the microscopic neural network architecture, the model structure is optimized to reduce transmission overhead and improve training efficiency.
It effectively solves the bottleneck problem of data transmission, improves training speed and accuracy, reduces transmission overhead, and maintains the feasibility and convenience of operations.
Smart Images
Figure CN120409628A_ABST
Abstract
Description
Technical Field
[0001] The present invention relates to the technical field of distributed machine learning, and particularly relates to a distributed model training method based on key layers of neural networks. Background Art
[0002] With the rapid development of artificial intelligence technology and the wide popularization of various intelligent devices, both the number of devices and the amount of data generated in the Internet of Things have increased significantly. Neural network models are a type of artificial intelligence technology widely used in intelligent devices. Generally, the model capabilities are positively correlated with their parameter scales. However, the huge computing power overhead required for their training is often unbearable for a single device. The distributed model training method can coordinate multiple devices to jointly calculate a certain task, and the use of parallel computing disperses the huge time and computing power consumption brought by centralized learning, and has gradually become the mainstream neural network model training method.
[0003] For the distributed model training method of centralized data parallelism, it includes several devices, which are divided into training nodes and aggregation nodes according to their functions. The total data set is divided into several non-overlapping subsets and assigned to each training node. The training is carried out in multiple rounds of iteration. In each round of training, each training node independently completes forward and backward propagation based on the selected model and generates parameter gradients. These gradients are then sent to the aggregation node for aggregation to generate new global model parameters and are sent back to all training nodes as the initial parameters for the next round of iteration.
[0004] The distributed model training method has a data transmission bottleneck. The throughput and bandwidth of the internal computing devices of nodes are often greater than the communication bandwidth between nodes, resulting in the data transmission speed becoming the bottleneck of the model training speed. The solutions are mainly divided into two categories: transmission-centered and computation-centered. The transmission-centered method reduces the total data transmission volume by selectively transmitting some gradient data, such as Top-K; the computation-centered method focuses on efficient computing inside the model and reduces the amount of data that needs to be transmitted from the source by training only some rather than all parameters.
[0005] In addition, using a general model structure for training on different data sets may lead to a mismatch between the model and the data set, thereby affecting the training efficiency and results. However, designing a dedicated model for each data set will bring high costs and lack practical feasibility. The existing solution is to adopt Neural Architecture Search (NAS), which mainly consists of three modules: search space, search strategy, and performance evaluation strategy. It comprehensively formulates the matching of the data set and the neural network architecture and the training of the neural network weights as a two-layer optimization problem, and attempts to use reinforcement learning or gradient methods to find the best architecture and optimal weights in the preset search space for a specific data set.
[0006] However, there is still a lack of a distributed model training method that can solve the above two problems at one time. Summary of the Invention
[0007] In view of this, the present invention provides a distributed model training method based on the key layers of a neural network. According to the determined key layers, parameter transmission from training nodes to a convergence node is completed, realizing low-transmission-overhead distributed training of the neural network.
[0008] A distributed model training method based on the key layers of a neural network provided by the present invention includes the following steps:
[0009] Deploy the model to be trained on a single machine, and use differentiable neural network search to determine the key layers in the model to be trained layer by layer;
[0010] Deploy the model to be trained to training nodes in a multi-machine distributed deployment manner. The training nodes use local data to calculate parameter gradients locally. The gradients of the key layers are cumulatively in error locally. According to the error accumulation results, select some gradients for sparse transmission to the convergence node; the convergence node aggregates the gradients from all training nodes, updates its own model parameters using the aggregated gradients, and transmits the updated model parameters back to each training node until the model converges to end the training process and complete the training of the model to be trained.
[0011] Further, the method of using differentiable neural network search to determine the key layers in the model to be trained layer by layer is as follows:
[0012] Determine the search space to perform optimal structure search layer by layer of the model, design a loss function, and use the differentiable neural network search method with minimizing the loss function as the optimization goal to find the optimal architecture in the search space. Stop training when reaching a preset number of rounds, and the obtained optimal architecture is the key layer.
[0013] Further, the method of determining the search space is: copy each layer of the model to be trained into two blocks with the same structure and initial parameters, including a parameter-adjustable block and a parameter-fixed block. The parameter-adjustable block dynamically updates parameters during training, while the parameter-fixed block keeps the initial parameters unchanged. All blocks of the model to be trained constitute the search space.
[0014] Further, the parameters of the model to be trained are initialized using a normal distribution.
[0015] Further, the method of using differentiable neural network search to determine the key layers in the model to be trained is as follows:
[0016] Step 1.1: Establish a loss function:
[0017]
[0018] in, To ensure the loss of accuracy, To control the total number of adjustable parameters, Num(out i1 ) is the number of parameters of the parameter-adjustable block outi1 of the i-th layer of the model to be trained, β is the weight coefficient; o is the structural parameter, o={c i1 ,c i2}, c i1 is the weight coefficient of the parameter-adjustable block in layer i, c i2 is the weight coefficient of the fixed block of parameters in layer i, and c i1 +c i2 =1;w o are model parameters;
[0019] Step 1.2: Use Gradient Update the parameter o, where ξ is the weight coefficient; use the gradient Update parameter w o ; is the gradient set of model parameters, is the training loss;
[0020] Step 1.3, iterate step 1.2 for multiple rounds until the preset training rounds are reached and the training process ends; select c i1 The k layers with larger values are taken as key layers.
[0021] Furthermore, the parameters o and w o The calculation method is:
[0022] Step 2.1: Introduce Gumbel noise gi1 for the parameter adjustable block and introduce Gumbel noise gi2 for the parameter fixed block. i1 , g i2 ~Gumbel(0,1), where sampling from the Gumbel distribution is computed by g=-log(-log(U)), and U ~uniform(0,1);
[0023] Step 2.2, calculate the original output value with noise:
[0024] z i1 =log(c i1 )+g i1 , z i2 =log(c i2 )+g i2
[0025] Among them, zi1 is the original output value of the parameter adjustable block, and zi2 is the original output value of the parameter fixed block;
[0026] Step 2.3: Use the Softmax function to calculate the weighted coefficients y of the output data of the parameter-adjustable block and the parameter-fixed block i1 and y i2 :
[0027]
[0028] where y i1 and y i2 satisfy y i1 + y i2 = 1;
[0029] Step 2.4: Aggregate the output data outi of the i-th layer:
[0030]
[0031] where x i is the input data of the i-th layer, outi1 is the parameter-adjustable block, and outi2 is the parameter-fixed block.
[0032] Furthermore, the method of accumulating errors locally for the gradients of the key layer and selecting some gradients for sparse transmission to the aggregation node is as follows:
[0033] Sum the absolute value of the gradient of the parameter p of the key layer at the t-th round with the cumulative gradient to obtain where has an initial value of 0 at the start of training; select the gradients of the k parameters with larger at the t-th round and transmit them to the aggregation node; after transmission, set the cumulative gradient of this parameter to zero.
[0034] Furthermore, the method of updating its own model parameters using the aggregated gradients is as follows:
[0035] Aggregate the gradients by taking the average of the gradients of N training nodes:
[0036]
[0037] where is the gradient set formed by the gradients of all parameters of the n-th training node;
[0038] Use the aggregated average gradient to update the model parameter wt of the aggregation node at the t-th round, and the update method is:
[0039]
[0040] Among them, η is the learning rate, and wt-1 is the model parameter of the aggregation node after the (t-1)-th round of update.
[0041] Beneficial effects:
[0042] By combining transmission-centered and computation-centered methods, the present invention uses the matching degree of data and model as the standard for selecting efficient parameters, adopts differentiable neural network architecture search technology, searches for the model architecture layer by layer, and defines the layer containing key parameters selected as the key layer; in subsequent distributed training, in order to further reduce the transmission overhead, gradient sparsification transmission and error accumulation strategy are adopted for the key layer to transmit parameters, which can effectively solve the mismatch problem between the model structure and training data, greatly alleviate the data transmission bottleneck problem, and at the same time retain the feasibility and convenience of the operation, showing advantages in inference accuracy and training speed in the comparison experiment with the existing scheme. Description of the drawings
[0043] Figure 1 It is a schematic diagram of the processing flow of a distributed model training method based on the key layer of a neural network provided by the present invention. Specific embodiments
[0044] The following examples are listed in conjunction with the drawings to describe the present invention in detail.
[0045] A distributed model training method based on the key layer of a neural network provided by the present invention has a processing process as Figure 1 shown, and specifically includes the following steps:
[0046] Step 1: Deploy the model to be trained on a single host, and use differentiable neural network architecture search to determine the key layer in the model to be trained.
[0047] Specifically, the method for determining the key layer using differentiable neural network architecture search is as follows: design the search space, perform optimal structure search layer by layer of the model; design the loss function, comprehensively consider the model accuracy and the number of adjustable parameters; use differentiable neural network architecture search technology, with minimizing the loss function as the optimization goal, search for the optimal architecture in the search space, and then select the key layer. By using the total data set to train on one machine and stopping training when reaching the preset number of rounds without the model converging, the key layer in the model to be trained can be determined.
[0048] Furthermore, the method for determining the key layer in the model to be trained provided by the present invention includes:
[0049] Step 1.1. Determine the search space: Copy each layer of the model to be trained into two blocks with the same structure. One block dynamically updates its parameters during training, and this block is denoted as the parameter-tunable block; the other block keeps its initial value unchanged, and this block is denoted as the parameter-fixed block; the search space is composed of all blocks of the model to be trained; the parameters of the model to be trained are initialized with a normal distribution, and the two blocks of each layer have the same initial parameters.
[0050] The definition of neural architecture search NAS is as follows: Suppose the search space is denoted as the dataset is D, denotes the validation loss, denotes the training loss, w o denotes the model parameters of architecture o. The goal of NAS is to find the best architecture with the optimal validation performance The mathematical representation is as follows:
[0051]
[0052] To balance the search effect and search time, different from the fine-grained search usually adopted by existing NAS methods, the present invention conducts model structure search in units of layers. Specifically, each layer of the selected model needs to be copied into two blocks with the same structure. One block dynamically updates its parameters during training, and the other block keeps its initial value unchanged. All blocks constitute the search space. The model parameters are initialized with a normal distribution, and the two blocks of each layer have the same initial parameters.
[0053] Furthermore, in order to evaluate the criticality of the two blocks in each layer, the present invention uses the Gumbel Softmax method for quantitative analysis, and let c i1 be the weight coefficient of the parameter-tunable block of the i-th layer, c i2 be the weight coefficient of the parameter-fixed block of the i-th layer, and c i1 + c i2 = 1, and both are initialized to 0.5, c i is the parameter of the i-th layer. c i1 reflects the necessity degree of a certain layer of the model being trained, so it can be used as a quantitative index of the critical degree.
[0054] The process of using the Gumbel Softmax method for quantitative analysis is as follows:
[0055] Step 1.1.1. Introduce Gumbel noise g i1 to the parameter-tunable block, and introduce Gumbel noise g i2 to the parameter-fixed block, g i1 , g i2~Gumbel(0, 1), where the sampling of the Gumbel distribution can be calculated by g = -log(-log(U)), and U ∼ Uniform(0, 1).
[0056] Step 1.1.2: Calculate the noisy original output value logits as shown in the following formula:
[0057] z i1 = log(c i1 ) + g i1 z i2 = log(c i2 ) + g i2
[0058] where z i1 is the original output value of the tunable parameter block, and z i1 is the original output value of the fixed parameter block.
[0059] Step 1.1.3: Use the Softmax function to calculate the weighting coefficients of the output data of the tunable parameter block and the fixed parameter block.
[0060] Use the temperature parameter τ to control the smoothness of the distribution. The calculation method of the output probability of Gumbel-Softmax is as follows:
[0061]
[0062] where y i1 and y i2 are the weighting coefficients for the two operations, satisfying y i1 + y i2 = 1.
[0063] Step 1.1.4: Aggregate the output data out of the i-th layer i , and the calculation method is as follows:
[0064]
[0065] where x i is the input data of the i-th layer, out i1 represents the calculation of the input data by the tunable parameter block, and out i2 represents the calculation of the input data by the fixed parameter block. After quantifying the importance index, the importance between multiple layers of the model can be compared according to the value of c i1 .
[0066] Step 1.2: Calculate the loss function.
[0067] The loss function established in the present invention consists of two parts, as shown in the following formula:
[0068]
[0069] Among them, is the loss term to ensure the accuracy of the task. Specifically, the cross-entropy loss function is adopted in the classification task; is the loss term to control the total number of adjustable parameters. Num(out i1 ) is the number of parameters in the adjustable parameter block of the i-th layer, and β is the weight coefficient of this term.
[0070] Step 1.3: Select the key layers.
[0071] In the present invention, on one device, using the total data set in the designed search space, based on the above loss function, the differentiable neural architecture search technology is adopted for model training. During the training process, an approximate iterative optimization step is used to alternately update two parameters o and w o . The specific process is as follows:
[0072] Step 1.3.1: Update the parameter o using the gradient . Among them, o is the structural parameter, o = {c i1 , c i2}; the value of ξ is usually set to be consistent with the learning rate of the model parameter w o .
[0073] Step 1.3.2: Update the parameter w using the gradient . o
[0074] Step 1.3.3: Conduct multiple rounds of iteration according to Step 1.3.1 to Step 1.3.2 until the preset number of training rounds is reached.
[0075] After the training is completed, according to the value of c i1 , the criticality of each layer block is sorted in descending order, and the first k layers are selected as the key layers.
[0076] Step 2: Deploy the model to be trained to each training node by using the multi-machine distributed deployment method. The training node calculates the parameter gradient locally using its own data, but does not perform parameter update; sparsely transmit the gradient of the key layer, and accumulate the error of the remaining gradients locally; the aggregation node aggregates the gradients transmitted from all training nodes and updates its own model parameters with the aggregated gradients; the aggregation node returns the parameters to each training node. This step is multi-machine distributed model training and needs to be trained until the model converges.
[0077] Step 2.1: Local gradient calculation.
[0078] Each training node saves a copy of the selected model and performs the same initialization using a normal distribution. Each training node calculates the gradient on the local data using the cross-entropy loss function
[0079] Step 2.2, Gradient Sparse Transmission and Error Accumulation.
[0080] After each training node completes one backpropagation, the Top-K method is used for a part of the gradients of the key layers to transmit them to the aggregation node, and the other part accumulates errors locally. That is, the gradient of parameter p in the t-th round is summed with the cumulative gradient to obtain where has an initial value of 0 at the start of training; select the gradients of the k parameters with larger in the t-th round and transmit them to the aggregation node; after transmission, set the cumulative gradient of this parameter to zero. Thus, error accumulation can ensure that sparse transmission converges at the same speed as standard transmission. For non-key layer gradients, no processing is done, that is, they are neither transmitted to the aggregation node nor accumulated locally.
[0081] Step 2.3, The aggregation node receives the gradients of all parameters sent by N training nodes The gradients of all parameters of the n-th training node form the gradient set of this training node The aggregation node performs gradient aggregation and parameter backpropagation, where N is the total number of training nodes.
[0082] The present invention aggregates gradients by taking the average of the gradients of N training nodes:
[0083]
[0084] Using the aggregated average gradient update the model parameter w of the aggregation node in the t-th round t , and the update method is:
[0085]
[0086] where η is the learning rate and w t-1 is the model parameter of the aggregation node after the update in the (t - 1)-th round. Transmit the updated parameter w t back to all training nodes as the initial model parameter for each training node in the next round of iteration.
[0087] Step 2.4, Conduct multiple rounds of iteration according to Steps 2.1 to 2.3 until the model converges or reaches the preset maximum number of training rounds.
[0088] Finally, it should be noted that: Obviously, those skilled in the art can make various modifications and variations to the present invention without departing from the spirit and scope of the present invention. Thus, if these modifications and variations of the present invention fall within the scope of the claims of the present invention and their equivalent technologies, the present invention also intends to include these modifications and variations.
[0089] The above is only an embodiment of the present invention, but it cannot be used to limit the scope of the present invention. Any structural changes made based on the present invention, as long as the essence of the present invention is not lost, should be regarded as falling within the protection scope of the present invention and being restricted. Those skilled in the art can clearly understand that for the convenience and brevity of description, the specific working process and related descriptions of the above-described method can refer to the corresponding process in the foregoing method embodiment, and will not be repeated here.
[0090] The term "comprising" or any other similar term is intended to cover non-exclusive inclusion, such that a process, method, article, or apparatus / method comprising a series of elements includes not only those elements but also other elements not expressly listed, or also includes elements inherent in these processes, methods, articles, or apparatuses / methods.
[0091] So far, the technical solutions of the present invention have been described by listing further embodiments. However, it is easy for those skilled in the art to understand that the protection scope of the present invention is obviously not limited to these specific embodiments. Without departing from the principle of the present invention, those skilled in the art can make equivalent changes or substitutions to the relevant technical features, and the technical solutions after these changes or substitutions will all fall within the protection scope of the present invention.
[0092] In summary, the above is only the preferred embodiment of the present invention and is not used to limit the protection scope of the present invention. Any modifications, equivalent substitutions, improvements, etc. made within the spirit and principle of the present invention shall be included within the protection scope of the present invention.
Claims
1. A distributed model training method based on key layers of a neural network, characterized in that, It includes the following steps: Deploy the model to be trained on a single machine, and use differentiable neural network search to determine the key layers in the model to be trained layer by layer; Deploy the model to be trained to the training nodes in a multi-machine distributed deployment manner. The training nodes calculate the parameter gradients locally using local data. The gradients of the key layers are accumulated locally for errors, and some gradients are selected for sparse transmission to the aggregation node according to the error accumulation results. The aggregation node aggregates the gradients from all training nodes, updates its own model parameters using the aggregated gradients, and transmits the updated model parameters back to each training node until the model converges to end the training process and complete the training of the model to be trained.
2. The distributed model training method according to claim 1, wherein The method of using differentiable neural network search to determine the key layers in the model to be trained layer by layer is as follows: Determine the search space to perform optimal structure search layer by layer of the model, design the loss function, and use the differentiable neural network search method with minimizing the loss function as the optimization goal to find the optimal architecture in the search space. Stop training when the preset number of rounds is reached, and the obtained optimal architecture is the key layer.
3. The distributed model training method according to claim 2, wherein The method of determining the search space is as follows: Copy each layer of the model to be trained into two blocks with the same structure and initial parameters, including a block with adjustable parameters and a block with fixed parameters. The block with adjustable parameters dynamically updates the parameters during training, while the block with fixed parameters keeps the initial parameters unchanged. All the blocks of the model to be trained constitute the search space.
4. The distributed model training method according to claim 3, wherein The parameters of the model to be trained are initialized with a normal distribution.
5. The distributed model training method according to claim 3, wherein The method of using differentiable neural network search to determine the key layers in the model to be trained is as follows: Step 1.1: Establish the loss function; Among them, is the loss term to ensure accuracy, is the loss term to control the total number of adjustable parameters, Num(out i1 ) is the adjustable parameter block out of the i-th layer of the model to be trained i1 The number of parameters, β is the weight coefficient; o is the structure parameter, o = {c i1 , c i2}, c i1 is the weight coefficient of the adjustable parameter block of the i-th layer, c i2 is the weight coefficient of the fixed parameter block of the i-th layer, and c i1 + c i2 = 1; w o is the model parameter; Step 1.2: Use the gradient to update the parameter o, where ξ is the weight coefficient; use the gradient to update the parameter w o ; is the gradient set of the model parameters, is the training loss; Step 1.3: Iteratively execute Step 1.2 for multiple rounds until a preset number of training rounds is reached, and end the training process; select k layers with larger c i1 values as the key layers.
6. The distributed model training method according to claim 3, wherein The parameters o and w o are calculated as follows: Step 2.1: Introduce Gumbel noise g into the parameter - adjustable block i1 , and introduce Gumbel noise g into the parameter - fixed block i2 , g i1 , g i2 ~Gumbel(0, 1), where the sampling of the Gumbel distribution is calculated by g = -log(-log(U)), U~Uniform(0, 1); Step 2.2: Calculate the raw output value with noise; z i1 = log(c i1 ) + g i1 , z i2 = log(c i2 ) + g i2 where z i1 is the original output value of the parameter-adjustable block, and z i1 is the original output value of the parameter-fixed block; Step 2.3: Use the Softmax function to calculate the weighted coefficients y of the output data of the parameter-adjustable block and the parameter-fixed block i1 and y i2 : where y i1 and y i2 satisfy y i1 + y i2 = 1; Step 2.4, aggregate the output data out of the i-th layer i : where x i is the input data of the i-th layer, out i1 is the parameter-adjustable block, out i2 is the parameter-fixed block.
7. The distributed model training method according to claim 1, wherein The method of accumulating the gradients of the key layers locally for errors and selecting some gradients for sparse transmission to the aggregation node according to the error accumulation results is as follows: The gradient of the parameter p of the key layer at the t-th round and the cumulative gradient are summed to obtain where the initial value at the start of training is 0; select the gradients of the k parameters with larger in the t-th round and transmit them to the aggregation node; after transmission, set the cumulative gradient of this parameter to zero.
8. The distributed model training method according to claim 1, wherein The method of updating its own model parameters using the aggregated gradients is as follows: Aggregate the gradients by taking the average of the gradients of N training nodes; Among them, is the gradient set formed by the gradients of all parameters of the nth training node; Using the average gradient obtained by aggregation Update the model parameter w of the aggregation node in the t-th round t , and the update method is as follows: where η is the learning rate, and w t-1 are the model parameters of the aggregation node after the (t - 1)-th round of update.