Model training method and system based on alternative adaptation momentum optimization

Abstract

The invention discloses a model training method and system based on alternative adaptive momentum optimization, and solves the technical problems that the memory efficiency of large-scale matrix optimization is low and training of a larger-scale model cannot be supported due to an existing model training method based on an adaptive optimization algorithm. The method comprises the steps that initial model parameters are acquired and marked as leaf nodes, and a gradient space used for storing first-order momentum is constructed; performing preprocessing according to the type of the initial model parameter, and outputting a target model parameter and model parameter analysis data; calculating unbiased first-order momentum of the current step number based on target model parameters and historical first-order momentum in a gradient space, and obtaining second-order momentum through element-by-element square; and performing odd-even alternate updating on the second-order momentum vector factor through an alternate adaptive momentum method, and iteratively updating a target model parameter according to an adaptive step length rule in combination with model parameter analysis data and unbiased first-order momentum to obtain a pre-training optimal parameter and construct a pre-training model.

Description

Technical Field

[0001] This invention relates to the field of artificial intelligence technology, and in particular to a model training method and system based on alternating adaptive momentum optimization. Background Technology

[0002] With the continuous iteration and breakthroughs in deep learning technology, large language models (LMs) based on the Transformer architecture have become a core technological pillar in fields such as natural language processing and artificial intelligence. Thanks to algorithm optimization and improved computing power, the parameter scale of these models has achieved exponential growth, gradually moving from the billion-level parameter scale in the early days to the trillion-level parameter scale. The model's representation ability and task adaptability have also been greatly improved, driving performance breakthroughs in various complex downstream tasks.

[0003] When training large language models with extremely large parameters, the dimensionality of the weight matrices of the core components within the model increases dramatically. Specifically, the Attention layer and FFN layer (Feed-Forward Network), key components of the Transformer architecture, contain linear projection matrices that serve as the core carriers of model parameters. Their dimensionality skyrockets in tandem with the overall increase in the number of parameters, directly posing a severe challenge to key metrics such as memory usage and computational complexity during model training. This becomes a crucial prerequisite restricting the efficient training of large models.

[0004] Existing model training methods based on adaptive optimization algorithms (such as Adam) must explicitly maintain a parameter matrix of the same size (m). The first-order momentum matrix (M) and second-order momentum matrix (U) are n). This mechanism introduces an additional memory overhead of up to 2mn. When training large models, the optimizer states often dominate the memory consumption (accounting for about 60%-70% of the total memory). The linear memory overhead directly leads to single-card memory overflow (OOM), severely limiting the expansion of model size and the choice of batch size, which is the core bottleneck restricting the efficient training of large models. For example, when fine-tuning all parameters of the GPT-3 model (Generative Pre-trained Transformer 3) with 175 billion (175B) parameters, the first and second momentum states maintained by the Adam optimizer alone require about 1.4TB of video memory space (based on FP32 (Floating-Point) storage computation). This far exceeds the storage limit of a single top-tier computing card (such as NVIDIA A10080GB), forcing the training task to be deployed on an expensive distributed cluster containing dozens of graphics cards. This requires extremely high computing power deployment costs and energy consumption, resulting in low memory efficiency for large-scale matrix optimization and making it impossible to support the training of larger-scale models. Summary of the Invention

[0005] This invention provides a model training method and system based on alternating adaptive momentum optimization, which solves the technical problem that existing model training methods based on adaptive optimization algorithms suffer from low memory efficiency in large-scale matrix optimization and cannot support the training of larger-scale models.

[0006] The first aspect of this invention provides a model training method based on alternating adaptive momentum optimization, comprising:

[0007] Obtain the initial model parameters and mark them as leaf nodes to construct a gradient space for storing first-order momentum;

[0008] Based on the model parameter type of the initial model parameters, the initial model parameters are preprocessed to output the target model parameters and model parameter parsing data;

[0009] Using a predefined loss function, unbiased first-order momentum is calculated based on the predefined training data distribution, the target model parameters, and the first-order momentum of the historical steps in the gradient space used to store first-order momentum, and the unbiased first-order momentum of the current step is output.

[0010] Perform an element-wise squaring operation on the unbiased first-order momentum of the current step and output the second-order momentum of the current step.

[0011] The alternating adaptive momentum method is used to update the second-order momentum vector factor of the current step number by alternating odd and even based on the analytical data of the model parameters, and the updated second-order momentum vector factor is determined.

[0012] An adaptive step size rule is used to iteratively update the target model parameters based on the model parameter parsing data, the updated second-order momentum vector factor, and the unbiased first-order momentum to obtain the pre-trained optimal parameters. A pre-trained model is then constructed based on the pre-trained optimal parameters.

[0013] Optionally, the preprocessing of the initial model parameters based on the model parameter type of the initial model parameters to output target model parameters and model parameter parsing data includes:

[0014] If the model parameter type of the initial model parameters is matrix, then the number of rows and columns of the initial model parameters are recorded to obtain the model parameter parsing data, and the initial model parameters are used as the target model parameters;

[0015] If the model parameter type of the initial model parameters is a high-dimensional tensor, then the dimensionality of the initial model parameters is recorded to obtain the model parameter parsing data, and the view of the initial model parameters is reshaped to obtain the target model parameters.

[0016] Optionally, the step of using a predefined loss function to calculate unbiased first-order momentum based on a predefined training data distribution, the target model parameters, and the first-order momentum of historical steps in the gradient space used to store first-order momentum, and outputting the unbiased first-order momentum of the current step, includes:

[0017] Randomly select training samples from the predefined training data distribution, use the training samples as input to the model corresponding to the target model parameters, and obtain the output result;

[0018] The loss value for the current step is calculated based on the output result according to the predefined loss function.

[0019] Based on the loss value at the current step, calculate the stochastic gradient of the loss value at the current step with respect to the target model parameters;

[0020] Perform in-situ scaling on the first-order momentum of the historical steps in the gradient space used to store the first-order momentum and the stochastic gradient to obtain the scaled first-order momentum of the historical steps and the scaled stochastic gradient.

[0021] The scaled stochastic gradient is accumulated into the scaled first-order momentum of the historical steps to obtain the first-order momentum of the current step.

[0022] The first-order momentum of the current step number is corrected for deviation to obtain the unbiased first-order momentum of the current step number.

[0023] Optionally, the vector factor includes row vector factors and column vector factors; the updated second-order momentum vector factor includes the row vector factor and column vector factor of the current step; the step of using the alternating adaptive momentum method to update the second-order momentum vector factor of the current step according to the model parameter analysis data by alternating odd and even values, and determining the updated second-order momentum vector factor, includes:

[0024] Determine if the current step number is the initial iteration step number;

[0025] If so, calculate the initial gradient norm mean based on the dimensional information of the parsed data of the model parameters;

[0026] Set each element value in the row vector factor and column vector factor of the second momentum of the current step number to the square root of the mean of the initial gradient norm to obtain the initial row vector factor and the initial column vector factor.

[0027] Use the initial row vector factor and the initial column vector factor as the corresponding row vector factor and column vector factor of the current step;

[0028] If not, determine whether the current step number is odd or even;

[0029] If the current step number is even, the column vector factor is fixed and used as the column vector factor of the current step number;

[0030] The optimal analytical solution for the row vector factors is obtained by sequentially performing matrix and vector multiplication, norm calculation, and division on the row vector factors.

[0031] An exponential moving average is performed on the optimal analytical solution of the row vector factor and the updated row vector factor of the historical steps to obtain the row vector factor of the current step.

[0032] If the current step number is odd, the row vector factor is fixed and used as the row vector factor for the current step number;

[0033] The optimal analytical solution for the column vector factors is obtained by sequentially performing matrix transpose, matrix-vector multiplication, norm calculation, and division on the column vector factors.

[0034] An exponential moving average is performed on the optimal analytical solution of the column vector factor and the updated column vector factor of the historical steps to obtain the column vector factor of the current step.

[0035] Optionally, the step-size adaptive rule iteratively updates the target model parameters based on the model parameter parsing data, the updated second-order momentum vector factor, and the unbiased first-order momentum to obtain the pre-trained optimal parameters, including:

[0036] Reconstruct the second-order moment matrix based on the updated second-order momentum vector factor;

[0037] The second-order moment matrix is ​​successively modified by order of magnitude and by deviation to obtain the modified second-order moment;

[0038] Adding a numerical stability constant to the modified second moment yields a modified second moment with a stability term;

[0039] Based on the unbiased first-order momentum and the modified second-order moment with a stabilizing term, the target model parameters are iteratively updated according to the adaptive step size rule to obtain the updated target model parameters.

[0040] When the iteration termination condition is reached, if the model parameter type of the updated target model parameters determined when the iteration termination condition is reached is a high-dimensional tensor, then the updated target model parameters determined when the iteration termination condition is reached are restored to the original tensor shape based on the model parameter parsing data. After verifying that the dimension is correct, they are determined as the pre-trained optimal parameters.

[0041] If the updated target model parameters determined when the iteration termination condition is reached are in matrix form, then the updated target model parameters determined when the iteration termination condition is reached will be used as the pre-trained optimal parameters.

[0042] Optionally, it also includes:

[0043] If the iteration termination condition is not met, the updated target model parameters are used as the new target model parameters, and the process jumps to the step of using a predefined loss function to calculate the unbiased first momentum based on the predefined training data distribution, the target model parameters, and the first momentum of the historical steps in the gradient space used to store the first momentum, and outputting the unbiased first momentum of the current step, until the iteration termination condition is met.

[0044] A second aspect of the present invention provides a model training system based on alternating adaptive momentum optimization, comprising:

[0045] The acquisition module is used to acquire initial model parameters, mark the initial model parameters as leaf nodes, and construct a gradient space for storing first-order momentum;

[0046] The preprocessing module is used to preprocess the initial model parameters based on the model parameter type of the initial model parameters, and output the target model parameters and model parameter parsing data;

[0047] The calculation module is used to perform unbiased first-order momentum calculation based on the predefined training data distribution, the target model parameters, and the first-order momentum of the historical steps in the gradient space used to store the first-order momentum using a predefined loss function, and output the unbiased first-order momentum of the current step.

[0048] The second-order momentum output module is used to perform element-wise squaring of the unbiased first-order momentum of the current step and output the second-order momentum of the current step.

[0049] The update module is used to perform odd-even alternating updates on the second-order momentum vector factor of the current step number based on the parsed data of the model parameters using the alternating adaptive momentum method, and to determine the updated second-order momentum vector factor.

[0050] The module is used to iteratively update the target model parameters using an adaptive step size rule based on the model parameter parsing data, the updated second-order momentum vector factor, and the unbiased first-order momentum, to obtain the pre-trained optimal parameters, and to construct a pre-trained model based on the pre-trained optimal parameters.

[0051] A third aspect of the present invention provides an electronic device, including a memory and a processor, wherein the memory stores a computer program, and when the computer program is executed by the processor, the processor performs the steps of the model training method based on alternating adaptive momentum optimization as described above.

[0052] The fourth aspect of the present invention provides a computer-readable storage medium having a computer program stored thereon, wherein the computer program, when executed, implements the model training method based on alternating adaptive momentum optimization as described above.

[0053] The fifth aspect of the present invention provides a computer program product comprising a computer program stored on a non-transitory computer-readable storage medium, the computer program comprising program instructions, wherein when the program instructions are executed by a computer, the computer performs the steps of the model training method based on alternating adaptive momentum optimization as described above.

[0054] As can be seen from the above technical solutions, the present invention has the following advantages:

[0055] The above-mentioned technical solution of the present invention provides a model training method based on alternating adaptive momentum optimization. The method involves obtaining initial model parameters and marking them as leaf nodes to construct a gradient space for storing first-order momentum; preprocessing the initial model parameters based on their type to output target model parameters and analytical model parameter data; using a predefined loss function, calculating unbiased first-order momentum based on a predefined training data distribution, target model parameters, and the first-order momentum of historical steps in the gradient space for storing first-order momentum, and outputting the unbiased first-order momentum of the current step; performing element-wise squaring on the unbiased first-order momentum of the current step to output the second-order momentum of the current step; and using the alternating adaptive momentum method to optimize the model training of the current step based on the analytical model parameter data. The second-order momentum vector factor is updated alternately using parity to determine the updated second-order momentum vector factor. An adaptive step-size rule is used to iteratively update the target model parameters based on the model parameter parsing data, the updated second-order momentum vector factor, and the unbiased first-order momentum, obtaining the pre-trained optimal parameters. A pre-trained model is then constructed based on these optimal parameters. Based on this scheme, this invention constructs a dedicated gradient space to store the first-order momentum, reusing the gradient storage characteristics of leaf nodes to avoid additional memory usage. The preprocessing stage adapts to subsequent algorithms according to parameter type, reducing redundant storage overhead. Furthermore, this invention employs an alternating minimization strategy to update the second-order momentum vector factor alternately using parity, eliminating the need to maintain a complete second-order momentum matrix of the same size as the parameter matrix, thus simplifying traditional adaptive optimization algorithms. The linear memory overhead is transformed into With its low-order overhead, it significantly improves the memory efficiency of large-scale matrix optimization and can support the training of larger-scale models. Attached Figure Description

[0056] To more clearly illustrate the technical solutions in the embodiments of the present invention or the prior art, the drawings used in the description of the embodiments or the prior art will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0057] Figure 1 The flowchart illustrates the steps of a model training method based on alternating adaptive momentum optimization provided in Embodiment 1 of the present invention.

[0058] Figure 2 This is a schematic diagram of the memory layout when running Adafactor, Adam, and Alada, as provided in Embodiment 1 of the present invention.

[0059] Figure 3 This is a structural block diagram of a model training system based on alternating adaptive momentum optimization provided in Embodiment 2 of the present invention. Detailed Implementation

[0060] This invention provides a model training method and system based on alternating adaptive momentum optimization, which relates to large-scale model training optimization technology in the field of artificial intelligence. It solves the technical problem that existing model training methods based on adaptive optimization algorithms suffer from low memory efficiency in large-scale matrix optimization, making it unable to support the training of larger-scale models.

[0061] To make the objectives, technical solutions, and advantages of the embodiments of the present invention clearer, the technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention. It should be noted that in the optional embodiments of the present invention, the object information and other related data involved require the permission or consent of the object when the embodiments of the present invention are applied to specific products or technologies, and the collection, use, and processing of related data must comply with the relevant laws, regulations, and standards of the relevant countries and regions. That is to say, if the embodiments of the present invention involve data related to the object, it needs to be obtained with the authorization and consent of the object, the authorization and consent of the relevant departments, and in compliance with the relevant laws, regulations, and standards of the country and region. If personal information is involved in the embodiments, the acquisition of all personal information requires the consent of the individual. If sensitive information is involved, the separate consent of the information subject is required, and the embodiments also need to be implemented with the authorization and consent of the object.

[0062] Terminology Explanation:

[0063] a) Alada: This refers to the Alternating Adaptation of Momentum Method proposed in this invention, used to reduce memory consumption in large-scale matrix optimization. This method performs a low-rank decomposition of second-order momentum based on Euclidean distance and employs an alternating update strategy to adapt to gradient changes. Simultaneously, it introduces a first-order momentum estimation mechanism that requires no additional memory consumption, targeting m... The optimization process for the n-matrix is ​​reduced from linear space complexity to sublinear complexity without significantly affecting optimization performance.

[0064] b) Memory-efficient large model fine-tuning: This refers to addressing the problem of excessive optimizer state memory consumption during fine-tuning of pre-trained large-scale deep learning models (such as BERT (Bidirectional Encoder Representations from Transformers), GPT (Generative Pre-trained Transformer), OPT (Open Pre-trained Transformer), and other Transformer architecture models) for specific downstream tasks. By decomposing the second-order momentum matrix into sublinear low-rank vectors and utilizing gradient accumulation to reuse storage space to maintain the first-order momentum, this invention reduces the additional memory overhead during fine-tuning from being proportional to the number of parameters. Significantly reduced to This allows for full parameter fine-tuning of large models with limited hardware resources.

[0065] c) Adaptive optimization methods: These refer to variants of Stochastic Gradient Descent (SGD) algorithms (such as Adam (Adaptive Moment Estimation), AdaGrad (Adaptive Gradient Algorithm), and RMSProp (Root Mean Square Propagation)). These methods automatically adjust the learning rate of each parameter by maintaining historical statistical information about the gradient (such as the first and second moments) to adapt to changes in the update frequency of different parameters and gradients.

[0066] d) First-order momentum: refers to the exponential moving average (EMA) of the gradient. During optimization, it is used to smooth random noise in the gradient and indicate the main direction of parameter updates. In this invention, the first-order momentum is denoted as... Furthermore, the gradient storage space is reused in the memory implementation.

[0067] e) Second-order momentum: refers to the exponential moving average (or an approximation) of the element-wise squared gradient. In adaptive methods, it is used to estimate the variance of the gradient, thereby scaling the learning rate element-wise (preprocessing) to address the ill-conditioned curvature problem. In this invention, the second-order momentum is denoted as... It is then reconstructed through the outer product of two low-dimensional vectors.

[0068] f) Rank-1 decomposition: This refers to the process of approximating a matrix as the outer product of two vectors (a column vector and a row vector). In this invention, it refers to decomposing the second-order moment matrix G of the gradient.2 Approximate decomposition into In the form of, and These are row vectors and column vectors, respectively.

[0069] g) Alternating Minimization: This refers to an iterative strategy for solving multivariate optimization problems. When the objective function is jointly non-convex with respect to all variables but convex with respect to a single variable block, the optimal solution is approximated by alternately fixing other variables and optimizing the current variable. This invention utilizes this strategy to alternately update factors. and To minimize the decomposition error.

[0070] h) Sublinear memory complexity: This refers to the rate at which the additional memory space required by the algorithm (excluding the space needed to store the model parameters themselves) grows less linearly with respect to the number of model parameters. In this invention, it specifically refers to reducing the optimizer's memory overhead from being proportional to the number of parameters. Reduced to a value proportional to the sum of the dimensions of the parameter matrices. .

[0071] i) Automatic Differentiation: This refers to a technique in computer programs that automatically calculates the derivative of a function by constructing a computational graph. Modern deep learning frameworks (such as PyTorch) widely employ this technique. This invention utilizes the gradient accumulation property of leaf nodes in the AD (Automatic Differentiation) framework to achieve first-order momentum updates with zero additional memory.

[0072] Please see Figure 1 , Figure 1 The flowchart illustrates the steps of a model training method based on alternating adaptive momentum optimization, as provided in Embodiment 1 of the present invention.

[0073] This invention provides a model training method based on alternating adaptive momentum optimization, comprising:

[0074] Step 101: Obtain the initial model parameters and mark them as leaf nodes to construct a gradient space for storing first-order momentum.

[0075] Initial model parameters refer to the basic parameters loaded during the initial stage of model training for subsequent iterative optimization. They cover the core optimizable parameters such as the weight matrices of the Attention layer and the Feedforward Neural Network (FFN) layer in the Transformer architecture, as well as various bias terms. They are the basic carriers for the model's representational capabilities and the core objects for gradient calculation and parameter updates.

[0076] Leaf nodes refer to parameter nodes in the computation graph of the automatic differentiation (AD) framework that have no preceding dependent nodes. After marking the initial model parameters as leaf nodes, these parameters can serve as the endpoints of gradient propagation. This can adapt to the gradient accumulation characteristics of the AD framework, avoid gradient storage redundancy caused by parameter correlation, and ensure the efficiency of gradient calculation and storage.

[0077] First-order momentum refers to the cumulative average of the gradients in each iteration during model training. It reflects the trend and direction of gradient changes and is one of the core bases for updating model parameters in adaptive optimization algorithms. It is used to balance the stability and convergence speed of gradient updates and improve the rationality of parameter updates.

[0078] Gradient space refers to a dedicated storage area for storing first-order momentum data. Its dimensions correspond one-to-one with the dimensions of the initial model parameters. It is used only to store first-order momentum information generated during the iteration process, without allocating additional redundant storage space, thus achieving efficient storage, retrieval, and updating of first-order momentum.

[0079] It should be noted that this invention provides a memory-optimized underlying architecture, aiming to solve the problem of additional storage overhead that typically arises from introducing first-order momentum estimation. In existing deep learning automatic differentiation frameworks (such as PyTorch), leaf nodes of the computation graph are usually allocated dedicated storage attributes (e.g., .grad) to store gradient information computed during backpropagation. To achieve the goal of "zero additional memory," this invention constructs a storage mechanism for gradient and momentum reuse, such as... Figure 2 As shown in (b) of the diagram. The core of this architecture lies in breaking the traditional practice of clearing gradients to zero before each iteration in the training process, and instead directly reusing the gradient attribute space of the leaf nodes to persistently store the first-order momentum. In specific implementation, the system does not have a first-order momentum M. t Instead of allocating a separate video memory buffer, the system directly stores the momentum value in the gradient attribute space. Before performing backpropagation, the system first performs an in-place scaling operation on the momentum value stored in this space from the previous time step, multiplying it by a decay factor. Subsequently, the system calculates the objective function loss value and applies it. The scaling factor. The scaling gradient is calculated when the automatic differentiation engine performs backpropagation. Based on the gradient accumulation characteristic of the framework, the values ​​will be automatically added to the aforementioned gradient attribute space. At this point, the values ​​in that space will naturally be updated. That is, the first-order momentum M required by the algorithm at the current moment. t +1. Through this design, the present invention retains the complete momentum properties while avoiding problems like those encountered with the Adam algorithm (e.g., Figure 2As shown in (a) in the figure, a separate first-order moment storage area is created, thereby reducing the optimizer's memory usage to a sublinear level comparable to momentum-free algorithms (such as Adafactor).

[0080] In this embodiment, the acquired initial model parameters cover the core optimizable parameters in the Transformer architecture, including the weight matrices of the Attention layer and the Feedforward Neural Network (FFN) layer, as well as various bias terms, ensuring coverage of all parameter dimensions required for model training. Marking the initial model parameters as leaf nodes leverages the gradient accumulation characteristic of the Automatic Differentiation (AD) framework, making the initial model parameters the endpoint of gradient propagation and avoiding gradient storage redundancy caused by pre-dependent nodes. Simultaneously, based on the dimensional information of the initial model parameters, a gradient space is constructed that matches each parameter dimension one-to-one. This gradient space exclusively stores the first-order momentum data generated during iteration, without allocating additional independent storage areas, achieving precise adaptation between first-order momentum storage and the parameter space. This invention, by accurately acquiring core initial model parameters, marking leaf nodes, and constructing a dedicated gradient space, reduces redundant memory overhead in initial parameter storage and first-order momentum storage from the source, lowering the GPU memory usage of the optimizer state and providing a foundation for further reducing memory consumption and improving memory efficiency.

[0081] Step 102: Preprocess the initial model parameters based on the model parameter type of the initial model parameters, and output the target model parameters and model parameter parsing data.

[0082] It should be noted that if the initial model parameters are in matrix form, the number of rows and columns are recorded as model parameter parsing data, and the initial model parameters are directly used as the target model parameters; if they are high-dimensional tensors, the size of each dimension is recorded as model parameter parsing data, and the view is reshaped to obtain the target model parameters, ensuring that the target model parameters are compatible with the subsequent momentum calculation and parameter update process.

[0083] Furthermore, step 102 may include the following sub-steps:

[0084] S21. If the model parameter type of the initial model parameter is matrix, then record the number of rows and columns of the initial model parameter to obtain the model parameter parsing data, and use the initial model parameter as the target model parameter.

[0085] S22. If the model parameter type of the initial model parameter is a high-dimensional tensor, then record the dimension size of the initial model parameter, obtain the model parameter parsing data, and reshape the view of the initial model parameter to obtain the target model parameter.

[0086] Matrix form refers to the two-dimensional shape category of the initial model parameters, consisting of a fixed number of rows and columns. It is commonly found in the linear projection matrix of the Transformer architecture. Its dimensional features can be fully represented by the number of rows and columns, and it can be adapted to the subsequent algorithm process without additional shape adjustments.

[0087] High-dimensional tensors refer to the multidimensional (dimension ≥ 3) morphological categories of initial model parameters. Their dimensional characteristics need to be represented by the specific size of each dimension to adapt to the complex feature storage requirements of the model. They need to be reshaped through view to adapt to the subsequent momentum calculation and parameter update process.

[0088] Model parameter parsing data refers to the initial model parameter dimensional information recorded during preprocessing. Matrix-form parameters correspond to the number of rows and columns, while high-dimensional tensors correspond to the size of each dimension, providing accurate dimensional support for subsequent operations such as second-order momentum factor updates and parameter shape restoration.

[0089] It should be noted that this invention primarily targets matrix variables (in matrix form). For high-dimensional tensors commonly used in fine-tuning of large models (such as convolution kernels), this invention does not address them. This invention proposes a general extension scheme based on ViewReshape. To apply the matrix factorization-based algorithm described above, the tensor Y needs to be reshaped into a two-dimensional matrix. To maximize memory savings (i.e., minimize m+n), the length and width of the reshaped matrix must be as close as possible. This invention employs the following strategy:

[0090] a) Find a split dimension index This makes the former Product of each dimension Product with the remaining dimensions The difference between them is the smallest, that is ≈ .

[0091] b) Utilize the .view() mechanism of modern deep learning frameworks (such as PyTorch) to logically treat tensors as m×n matrices without performing physical data copying.

[0092] c) Apply the Alada algorithm to the logical matrix.

[0093] This approach ensures This makes the extra memory used to store the p and q vectors negligible relative to the total video memory of the parameters.

[0094] Specifically, if the initial model parameters are in matrix form, the number of rows and columns of the initial model parameters in matrix form is accurately recorded. This ensures that the dimensional information of this type of parameter is completely preserved in the parsed model parameter data, providing accurate dimensional basis for subsequent operations such as second-order momentum vector factor updates and second-order moment matrix reconstruction. Simultaneously, the initial model parameters in matrix form can be directly reused as target model parameters, eliminating the need for additional shape adjustment operations and avoiding redundant memory usage caused by parameter form conversion. This invention reduces unnecessary memory overhead in the preprocessing stage and lowers the GPU memory usage ratio in the optimizer state by directly reusing matrix parameters and accurately recording dimensional information.

[0095] If the initial model parameters are of high-dimensional tensor type, the specific size of each dimension of the high-dimensional tensor is recorded in detail to form complete model parameter parsing data. This ensures the accuracy of parameter shape restoration after subsequent iterations. Simultaneously, a view reshaping operation is performed on the high-dimensional tensor, adjusting only the parameter's form display without changing its internal data information. This allows the reshaped target model parameters to adapt to subsequent processes such as unbiased first-order momentum calculation and alternating updates of second-order momentum vector factors. This invention adapts to subsequent algorithm requirements through view reshaping while avoiding redundant memory consumption caused by storing the original form of the high-dimensional tensor. Combined with dimensional information recording, it ensures the standardization of parameter processing, further reducing memory overhead and improving memory efficiency for large-scale matrix optimization.

[0096] Step 103: Using a predefined loss function, calculate the unbiased first-order momentum based on the predefined training data distribution, target model parameters, and the first-order momentum of the historical steps in the gradient space used to store the first-order momentum, and output the unbiased first-order momentum of the current step.

[0097] It should be noted that samples are drawn from a predefined training data distribution, input into the model corresponding to the target model parameters to obtain the output result, the loss value is calculated through a predefined loss function, and the unbiased first-order momentum of the current step is calculated according to the adaptive optimization rule by combining the historical first-order momentum stored in the gradient space, thus ensuring the rationality and accuracy of the momentum calculation.

[0098] Furthermore, step 103 may include the following sub-steps:

[0099] S31. Randomly select training samples from the predefined training data distribution, use the training samples as the input of the model corresponding to the target model parameters, and obtain the output results.

[0100] S32. Calculate the loss value for the current step based on the output result using a predefined loss function;

[0101] S33. Based on the loss value at the current step, calculate the stochastic gradient of the loss value at the current step with respect to the target model parameters;

[0102] S34. Perform in-situ scaling on the first-order momentum and stochastic gradient of the historical steps in the gradient space used to store the first-order momentum, to obtain the scaled first-order momentum and scaled stochastic gradient of the historical steps.

[0103] S35. Accumulate the scaled stochastic gradient into the scaled first-order momentum of the historical steps to obtain the first-order momentum of the current step.

[0104] S36. Correct the deviation of the first momentum of the current step to obtain the unbiased first momentum of the current step.

[0105] Predefined training data distribution refers to a pre-defined set of data and its distribution characteristics that conform to the model training task scenario. It is the source of training samples, ensuring that the extracted samples can meet the model training requirements and provide effective data support for gradient calculation.

[0106] Training samples refer to single or batch data randomly drawn from a predefined training data distribution for model training. They include input data and corresponding real labels and are the core data carriers for model learning and gradient calculation.

[0107] The model corresponding to the target model parameters refers to the complete model built around the preprocessed target model parameters. It can receive training sample inputs and output prediction results, and is the execution carrier for loss value and gradient calculation.

[0108] The output result refers to the prediction result output by the model after the training samples are input into the target model parameters. It is used to compare with the true labels of the samples to calculate the loss value and reflect the current learning effect of the model.

[0109] Predefined loss function: refers to a function that is set in advance to measure the degree of difference between the model output and the true label of the sample. Specifically, loss functions adapted to large model training, such as cross-entropy loss function and negative log-likelihood loss function, can be used.

[0110] It should be noted that the training samples for the model training task are extracted to ensure that the sample distribution is consistent with the predefined distribution. After the model is constructed by inputting the target model parameters, the model's prediction results (output results) are output. Based on the predefined loss function, the loss value for the current step is calculated according to the output results. The difference between the model's predicted output results and the true labels of the samples is quantified by the loss function to obtain the loss value for the current step, reflecting the current fitting effect of the model. Based on the loss value for the current step, the stochastic gradient of the loss value with respect to the target model parameters is calculated. The partial derivatives of the loss value with respect to the target model parameters are solved using the automatic differentiation (AD) framework to obtain the stochastic gradient, which characterizes the direction and magnitude of parameter adjustment. For storing the first-order... The first-order momentum and stochastic gradient of the historical steps in the momentum gradient space are scaled in-situ to obtain scaled first-order momentum and scaled stochastic gradient of the historical steps. The scaling operation is performed directly on the historical first-order momentum and stochastic gradient within the gradient space without allocating additional storage areas, avoiding redundant memory usage. The scaled stochastic gradient is accumulated into the scaled first-order momentum of the historical steps to obtain the first-order momentum of the current step. The accumulation operation is performed within the same gradient space, reusing storage space for efficient updating of the first-order momentum. Deviation correction is applied to the first-order momentum of the current step to obtain the unbiased first-order momentum of the current step. Deviation correction eliminates the influence of momentum deviation in the initial iteration stage, ensuring the accuracy of the momentum data. In this step, the entire unbiased first-order momentum calculation process is performed in-situ based on the existing gradient space, without needing to allocate additional independent storage areas, minimizing memory overhead during the first-order momentum calculation process, reducing the GPU memory usage ratio of the optimizer state, and effectively improving memory utilization efficiency.

[0111] Step 104: Perform element-wise squaring on the unbiased first-order momentum of the current step and output the second-order momentum of the current step.

[0112] It should be noted that the squaring operation is performed on each element of the unbiased first-order momentum at the current step. This does not change the parameter dimension structure; it only strengthens the gradient magnitude feature through squaring, transforming the gradient change trend reflected by the unbiased first-order momentum into intermediate data characterizing the gradient squared change law. This outputs the second-order momentum at the current step, providing core input for subsequent updates to the second-order momentum vector factor. This invention obtains second-order momentum through a simple and efficient element-wise squaring operation, eliminating the need to construct a complete second-order momentum matrix. This avoids the large memory consumption caused by storing a complete second-order momentum matrix in traditional methods, further reducing the memory overhead of the optimizer state and improving memory utilization efficiency.

[0113] Step 105: Using the alternating adaptive momentum method, update the second-order momentum vector factor of the current step number alternately with odd and even based on the analytical data of the model parameters, and determine the updated second-order momentum vector factor.

[0114] Vector factors include row vector factors and column vector factors.

[0115] It should be noted that this invention proposes a memory-efficient matrix optimization method called Alada (Alternating Adaptation of Momentum). This method is mainly used to address the problem of excessive memory consumption faced by Stochastic Gradient Descent (SGD) and its adaptive variants (such as Adam) when training large-scale deep learning models. This invention uses an alternating minimization strategy based on Euclidean distance to perform low-rank decomposition of the second-order gradient estimation, and combines it with a first-order momentum estimation mechanism with zero additional memory, thereby achieving a reduction in memory complexity from zero while maintaining the convergence and robustness of the algorithm. Reduce to .

[0116] Specifically, step 105 may include the following sub-steps:

[0117] S51. Determine if the current step number is the initial iteration step number;

[0118] S52. If so, calculate the mean of the initial gradient norm based on the dimensional information of the data parsed from the model parameters;

[0119] S53. Set the values ​​of each element in the row vector factor and column vector factor of the second momentum of the current step to the square root of the mean of the initial gradient norm to obtain the initial row vector factor and the initial column vector factor.

[0120] S54. Use the initial row vector factor and the initial column vector factor as the corresponding row vector factor and column vector factor of the current step.

[0121] S55. If not, determine whether the current step number is odd or even.

[0122] S56. If the current step number is even, the column vector factor is fixed and used as the column vector factor of the current step number.

[0123] S57. Perform matrix and vector multiplication, norm calculation and division on the row vector factors in sequence to obtain the optimal analytical solution of the row vector factors;

[0124] S58. Perform an exponential moving average on the optimal analytical solution of the row vector factor and the updated row vector factor of the historical steps to obtain the row vector factor of the current step.

[0125] S59. If the current step number is odd, fix the row vector factor and use it as the row vector factor for the current step number.

[0126] S510. Perform matrix transpose, matrix-vector multiplication, norm calculation and division on the column vector factors in sequence to obtain the optimal analytical solution of the column vector factors;

[0127] S511. Perform an exponential moving average on the optimal analytical solution of the column vector factors and the updated column vector factors of the historical steps to obtain the column vector factors of the current step.

[0128] The initial iteration step refers to the first iteration step in the model training process (corresponding to the initial time t=0 in the algorithm). It is the node that triggers the second-order momentum factor vector initialization operation and determines whether to execute the logic of initial gradient norm mean calculation and factor initialization.

[0129] It should be noted that in traditional adaptive methods (such as Adam), the second-order moment estimate Ut+1 is a matrix with the same dimension as the parameter matrix. The matrix is ​​updated using the following formula: This leads to a huge memory overhead. This invention employs the idea of ​​low-rank decomposition, assuming the second moment of the gradient... It can be approximated as the outer product of two vectors, i.e. ,in , Unlike existing techniques (such as Adafactor) that use generalized Kullback-Leibler divergence, this invention proposes using Euclidean distance as a metric for decomposition error. This is because... and It may not lie on the simplex, and the Euclidean distance has symmetry, making its physical meaning clearer. The goal of the decomposition is to minimize the following error:

[0130] ;

[0131] in, The second moment of the gradient at the current iteration step t is obtained by performing an element-wise squaring operation on the unbiased first momentum at the current step. The dimensional matrix represents the trend of the gradient squared at this step. It is the complete second-order moment matrix that needs to be maintained in traditional adaptive methods (such as Adam), and it is also the target matrix of the low-rank decomposition in this invention. is the row vector factor in the low-rank decomposition; These are the column vector factors in low-rank decomposition; The optimization objective of the low-rank decomposition in this invention is to calculate the values ​​of the row vector factor p and the column vector factor q so that the subsequent squared Euclidean distance (decomposition error) is minimized, and to determine the optimal second-order momentum factor vector. The second moment of the gradient outer product matrix The square of the Euclidean distance between them is an indicator used in this invention to measure the error of low-rank decomposition, and its physical meaning is the degree of approximate deviation between the two.

[0132] Directly solving the global optimal solution to the above problem typically requires Singular Value Decomposition (SVD), which is computationally expensive and difficult to perform in real-time during online training. Therefore, this invention proposes an imprecise alternating minimization strategy. Since when fixed... When p is constant, the objective function is strongly convex; conversely, it is not. Therefore, the optimal solution can be approximated by alternately fixing one variable and updating another, thus avoiding complex SVD calculations. Specifically, this invention alternately updates the analytical solutions of the following two subproblems during the iteration process:

[0133] fixed Update p:

[0134] ;

[0135] Fixed p, update :

[0136] ;

[0137] in, The optimal analytical solution (optimal analytical solution of row vector factors) for the subproblem of "fixing q and updating p"; This is the optimal analytical solution (optimal analytical solution of column vector factors) for the subproblem of "fixing p and updating q".

[0138] In the actual optimization loop, to smooth the estimation and further reduce the computational load, this invention does not completely recalculate p and q at every step. Instead, it treats the above analytical solution as a "target value" and updates it slowly using an exponential moving average (EMA). The update logic is related to the parity of the time step t: when t is even, q is kept constant. t No change, only update p t When t is odd, keep p. t No change, only q is updated t Through this alternating adaptation mechanism, what was originally needed... The stored second-order moment matrix is ​​compressed into two The vector.

[0139] Specifically, determine whether the current step is the initial iteration step; if so, combine the dimensional information of the model parameter parsing data and the stochastic gradient obtained in the first calculation to calculate the initial gradient norm mean. Here, the initial gradient norm mean corresponds to the gradient norm statistics at the initial moment, ensuring that the subsequent initialized second-order momentum factor vector matches the initial gradient magnitude; set each element value in the row vector factor and column vector factor of the second-order momentum of the current step to the square root of the initial gradient norm mean to obtain the initial row vector factor and initial column vector factor, thus ensuring the reasonable scale of the initial factors; use the initial row vector factor and initial column vector factor as the corresponding row vector factor and column vector factor of the current step to complete the configuration of the second-order momentum factor in the initial stage. If not, determine whether the current step number is odd or even. If the current step number is even, fix the column vector factor and use it as the column vector factor for the current step, that is, continue the column vector factor updated from the previous step without making any additional adjustments. Perform matrix and vector multiplication, norm calculation and division on the row vector factors in sequence to obtain the optimal analytical solution of the row vector factors. This calculation process is based on the operation of the second-order momentum estimation source and the current column vector factor, which is consistent with the analytical solution solution logic with the goal of minimizing the Euclidean distance. Perform an exponential moving average on the optimal analytical solution of the row vector factors and the row vector factors updated from the previous step to obtain the row vector factor for the current step. This operation follows the strategy in the algorithm of using the analytical solution as the target value and using an exponential moving average to achieve smooth updates. If the current step number is odd, the row vector factor is fixed and used as the row vector factor for the current step, meaning the row vector factor updated from previous steps is continued without additional adjustment. The column vector factors are then subjected to matrix transpose, matrix-vector multiplication, norm calculation, and division to obtain the optimal analytical solution for the column vector factors. This process is based on the operation between the second-order momentum estimation source and the current row vector factor, also aiming to minimize the Euclidean distance. An exponential moving average is then applied to the optimal analytical solution of the column vector factors and the updated column vector factors from previous steps to obtain the column vector factor for the current step. The exponential moving average ensures the smoothness of the column vector factor update. In this initial stage, the factor initialization using gradient norm mean matching ensures the reasonableness of the parameter scale. Subsequently, the alternating odd and even factor update strategy, combined with the smooth update method of exponential moving average, avoids complex full matrix operations and reduces the complexity required in traditional adaptive optimization algorithms. The stored second-order moment matrix is ​​compressed to only need The storage of two vector factors significantly reduces the memory overhead associated with second-order momentum and improves the memory efficiency of large-scale matrix optimization.

[0140] Step 106: Using an adaptive step size rule, the target model parameters are iteratively updated based on the model parameter parsing data, the updated second-order momentum vector factor, and the unbiased first-order momentum to obtain the pre-trained optimal parameters. Based on the pre-trained optimal parameters, a pre-trained model is constructed.

[0141] A pre-trained model refers to a complete model built based on the optimal parameters from pre-training. Specifically, it can be a large pre-trained model based on the Transformer architecture, such as the GPT series or Llama series. It can perform tasks related to natural language processing and is the final output of the pre-training process.

[0142] It should be noted that, based on the model parameter analysis data, the dimensional adaptation relationship of the target model parameters is clarified. The updated second-order momentum vector factors are reconstructed by performing an outer product operation to obtain the second-order moment estimate. Then, combined with the unbiased first-order momentum, the target model parameters are adjusted according to the element-by-element update logic of the adaptive step size rule. The iteration continues until the preset iteration termination condition is reached to obtain the pre-trained optimal parameters. Finally, a pre-trained model with the corresponding architecture is built based on the pre-trained optimal parameters.

[0143] Furthermore, step 106 may include the following sub-steps:

[0144] S61. Reconstruct the second-order moment matrix based on the updated second-order momentum vector factor;

[0145] S62. Perform order of magnitude correction and deviation correction on the second-order moment matrix in sequence to obtain the corrected second-order moment;

[0146] S63. Add a numerical stability constant to the modified second moment to obtain the modified second moment with a stability term;

[0147] S64. Based on the unbiased first-order momentum and the corrected second-order moment with a stabilizing term, the target model parameters are iteratively updated according to the adaptive step size rule to obtain the updated target model parameters.

[0148] S65. When the iteration termination condition is reached, if the model parameter type of the updated target model parameter determined when the iteration termination condition is reached is a high-dimensional tensor, then the updated target model parameter determined when the iteration termination condition is reached will be restored to the original tensor shape based on the model parameter parsing data. After verifying that the dimension is correct, it will be determined as the pre-training optimal parameter.

[0149] S66. If the updated target model parameters determined when the iteration termination condition is reached are in matrix form, then the updated target model parameters determined when the iteration termination condition is reached will be used as the optimal pre-training parameters.

[0150] The numerical stability constant refers to the small constant added to the corrected second moment. Its core function is to avoid the situation where the denominator is zero or the numerical overflow occurs in the subsequent adaptive step size calculation, thereby improving the robustness and stability of the parameter update process.

[0151] It should be noted that the Alada algorithm proposed in this invention (as shown in Table 1) aims to complete large-scale matrix optimization with low memory consumption through an alternating update strategy. The algorithm first initializes the first-order momentum matrix. Given an all-zero matrix, and specify the corresponding step size sequence { }, Attenuation parameters , and numerical stability constant In each iteration t, the algorithm first draws samples from the data distribution D. And calculate the stochastic gradient G t Subsequently, the first-order momentum is updated using an exponential moving average. And by correcting for the deviation, we obtain To optimize memory usage, the algorithm directly calculates the element-wise square of the corrected first-order momentum. It serves as an estimate of the second-order momentum (the second-order momentum at the current step), rather than storing the square of the original gradient. If the current time is the initial step, i.e., the initial iteration step (t=0), the algorithm will use the norm mean of the initial gradient. Initialize the second-order momentum factor vector and This ensures initial magnitude matching. The core second-order momentum update employs an alternating minimization strategy. Specifically, when the iteration number t is even, the algorithm keeps the vector q unchanged ( ← ), and estimate the source based on second-order momentum. Update vector Conversely, when t is odd, the algorithm keeps vector p unchanged. ← ), and update the vector This alternation mechanism approximates the second-order moment information by minimizing the Euclidean distance, avoiding complex full matrix operations. Finally, the algorithm uses outer product... Reconstructing the second moment estimate and combined with initialization items Obtained by deviation correction This is to eliminate initial bias. Finally, the corrected first-order momentum is used. and second momentum The target model parameters X are adjusted according to the adaptive step size rule. t Perform element-wise updates to obtain the updated target model parameters X. t+1 .

[0152] Table 1 Alada Algorithm

[0153]

[0154] Specifically, the updated row vector factors and column vector factors are multiplied by an outer product to obtain a second-order moment matrix that matches the dimension of the target model parameters, thus realizing the transformation from low-rank factors to second-order moment information. The second-order moment matrix is ​​then subjected to magnitude correction and bias correction to obtain the corrected second-order moment. The magnitude correction is used to match the magnitude of the initial gradient, while the bias correction combines initialization terms related to the norm and mean of the initial gradient to eliminate the bias effects of the initial iterations. A numerical stability constant is added to the corrected second-order moment to obtain a corrected second-order moment with a stability term, which is used to avoid numerical anomalies in subsequent calculations. Based on the unbiased first-order momentum and the corrected second-order moment with a stability term, the target model parameters are iteratively updated according to the adaptive step size rule to obtain the updated target model parameters. The step size is calculated by dividing the unbiased first-order momentum by the square root of the corrected second-order moment with a stabilizing term on an element-by-element basis, and then multiplying by the step size adjustment parameter. When the iteration termination condition is met (including the number of iterations reaching the preset maximum training steps, or the change in the loss value over several consecutive steps being less than the preset convergence threshold), if the updated target model parameters determined when the iteration termination condition is met are of the high-dimensional tensor type, then the model parameters are restored to their original tensor shape based on the model parameter parsing data (which records the original high-dimensional tensor dimensions). After verifying that the dimensions are correct, the updated target model parameters are determined as the optimal pre-training parameters. If the model parameters are of matrix form, then the updated target model parameters are directly used as the optimal pre-training parameters, since the matrix form parameters do not change shape during iteration and do not require additional adjustment. This invention reconstructs the second-order moment matrix through the outer product of low-rank factors, avoiding the high memory overhead of traditional complete second-order moment matrices. It combines correction and stabilization terms to ensure the accuracy and robustness of parameter updates, and obtains pre-trained optimal parameters based on parameter type differentiation. This continuously reduces memory consumption while improving the memory efficiency of large-scale matrix optimization, effectively alleviating the technical problem that existing adaptive optimization algorithms cannot support larger-scale model training due to low memory efficiency.

[0155] Optionally, it also includes:

[0156] If the iteration termination condition is not met, the updated target model parameters are used as the new target model parameters, and the process jumps to the step of using a predefined loss function to calculate the unbiased first momentum based on the predefined training data distribution, the target model parameters, and the first momentum of the historical steps in the gradient space used to store the first momentum, and outputs the unbiased first momentum of the current step, until the iteration termination condition is met.

[0157] It should be noted that when the iteration termination condition is not met, the updated target model parameters are used as the new target model parameters to maintain the continuity of parameter updates and the closed loop of the iteration process. Then, the process jumps to step S31 and continues to iterate until the iteration termination condition is met, thus ensuring the integrity of the parameter optimization process.

[0158] For comparison of technical effectiveness, existing technologies can be referenced. To train large models with limited resources, current memory-efficient optimization algorithms (such as Adafactor) are often forced to abandon explicit maintenance of first momentum to achieve sublinear memory usage. However, in the complex loss surface of large models, first momentum plays a crucial role in smoothing gradient noise and accelerating escape from local minima. Its absence causes such algorithms to converge significantly slower than the standard Adam algorithm when training large models, and they are prone to training instability (such as the "loss spike" phenomenon where the loss value suddenly and abnormally increases), making it difficult to balance "low memory usage" and "high training performance".

[0159] Existing second-order momentum estimation methods based on matrix factorization (such as Adafactor) typically use KL divergence as a metric for decomposition error. However, while KL divergence is mathematically used to measure the difference between two probability distributions (i.e., requiring numerical normalization), the second-order moment matrix of the gradient is merely a measure of numerical strength and does not possess the characteristics of a probability distribution. Therefore, using KL divergence in this scenario lacks clear physical meaning. Furthermore, KL divergence is asymmetric and not an ideal metric for measuring matrix approximation error (in contrast, Euclidean distance is more symmetric and intuitive). Simultaneously, in pursuit of computational efficiency, existing methods often struggle to obtain an exact solution to the decomposition problem, leading to insufficient accuracy in approximating the second-order gradient, ultimately affecting the update quality of large model training.

[0160] In recent years, large models have become a core technology in fields such as natural language processing and computer vision. These models (such as the GPT (Generative Pre-trained Transformer) series and the Llama (Large Language Model Meta AI) series) are typically based on the Transformer architecture and consist of hundreds or even thousands of stacked layers. Each layer contains a large number of linear projection operations, which essentially update a huge weight matrix. For example, in multi-head attention mechanisms and feed-forward networks, decision variables are represented as extremely large matrices. (Where m and n can be as high as tens of thousands). Therefore, the training process of a large model can be mathematically modeled as a stochastic optimization problem for a series of large-scale matrix variables:

[0161] (1)

[0162] in, Indicates a data sample. It is data distribution. The decision variables are in matrix form. Finding the minimum value of a matrix variable X of dimension m×n is the objective operation of the stochastic optimization problem corresponding to large model training, namely, finding the optimal weight matrix to minimize the objective function. This is the stochastic optimization objective function corresponding to the training of a large model. Its value is determined by the expected value of the samples in the data distribution. The core is to quantize the training objective (such as loss) corresponding to the model weight matrix X. For mathematical expectation operators, Based on decision variable X and data sample The objective function term. In a typical neural network training loop, updating the weight matrix of each layer is equivalent to solving a subproblem of the form of Equation (1).

[0163] Currently, the mainstream approach to solving such problems is adaptive optimization algorithms, with the Adam algorithm being the most representative. These methods accelerate convergence and improve training robustness by maintaining estimates of the first moment (momentum) and second moment (adaptive learning rate term) of the gradient. In the t-th iteration, Adam-like algorithms typically perform the following update steps:

[0164] (2)

[0165] in, It is the step size (learning rate). It is a very small constant (used for numerical stability). and These are the first-order moment estimate and the second-order moment estimate of the gradient, respectively. Division and square root operations are performed element-wise.

[0166] In the standard Adam algorithm, and It is updated using the Exponential Moving Average (EMA), and the specific formula is as follows:

[0167] (3)

[0168] in, It is the current stochastic gradient estimate (e.g., ), This represents the element-wise square of the gradient. , ∈[0,1) is the attenuation parameter. This term is mainly used to smooth gradient noise and improve the robustness of the optimization process; while The term is used to handle the ill-conditioning of the loss function terrain, achieving a second-order descent effect through element-wise preprocessing of the gradient. However, the main drawback of the standard Adam algorithm is its huge memory overhead. To perform the above update, the algorithm needs to explicitly maintain two matrices M and U with the same shape as the decision variable X. This means that for a matrix of size m... For a parameter matrix of n, the optimizer requires an additional 2^mn of memory. This is particularly problematic when training large-scale language models or processing extremely high-dimensional data. The linear memory overhead has become the main bottleneck limiting model size and training efficiency.

[0169] To address the memory bottleneck issue in adaptive optimization methods, an algorithm called Adafactor has been proposed in existing technologies. Adafactor is currently the closest low-memory adaptive optimization scheme to this invention and is widely used in the training of large language models.

[0170] a) Rank-One Factorization of Second Moments

[0171] The core idea of ​​Adafactor is to abandon the explicit storage of complete m. n-dimensional second-order moment matrix Instead of directly decomposing it into the outer product of two vectors, it approximates it as the second moment of the gradient. Specifically, it assumes the second moment of the gradient... It can be approximated as a row vector and column vectors The outer product, i.e. ≈ To find the optimal vectors p and q, Adafactor measures the decomposition error by minimizing the generalized Kullback-Leibler divergence (KL divergence). Its update rule is as follows:

[0172] (4)

[0173] In this way, Adafactor only needs to maintain two vectors p and q, and its memory complexity is reduced from that of Adam. Reduced to This saves memory space for mn.

[0174] b) Remove the first moment to further save memory.

[0175] To further reduce memory footprint, the authors of Adafactor suggest removing the first-moment estimation. (i.e., setting) =0). This means the algorithm directly uses the current stochastic gradient. To update parameters, instead of using smoothed momentum. This strategy saves the extra mn memory space used to store the M matrix. Therefore, the final parameter update formula for Adafactor is:

[0176] (5)

[0177] in, Represents from vector and The second-order moment estimate in the recovered (reconstructed) matrix form.

[0178] Based on the above, although adaptive optimization methods have been widely used in machine learning, existing solutions still have significant shortcomings when dealing with large-scale matrix optimization problems. The standard Adam algorithm requires explicitly maintaining two momentum matrices, the first moment (M) and the second moment (U), resulting in a memory overhead of 2^mn. The linear growth of memory at scale makes GPU memory a "hard constraint" for training large models. For example, when training a model with hundreds of billions of parameters, the optimizer state alone requires terabytes of GPU memory, forcing researchers to compromise by reducing model depth and width or significantly decreasing batch size, severely hindering the exploration of large model capabilities. Although Adafactor reduces GPU memory usage through low-rank decomposition, its design has the following shortcomings:

[0179] a) Removing first-order momentum leads to decreased convergence and robustness: In pursuit of ultimate memory compression, Adafactor typically removes the "first-order momentum" estimate (i.e., sets...). =0), directly using the current stochastic gradient for updating. While this approach saves GPU memory, it results in a lack of the smoothing effect of momentum during training, slowing down the algorithm's convergence speed and causing poor stability (robustness) when dealing with noisy data.

[0180] b) The physical meaning of the decomposition metric is ambiguous: Adafactor uses the generalized KL divergence to measure the second moment G. 2 With approximate matrix pq T The error between them. However, due to G 2 and pq T It may not lie on the simplex, and the KL divergence is asymmetric, which makes it ambiguous in a physical sense as a measure of decomposition error.

[0181] c) High computational cost: Exact decomposition problems based on KL divergence or Euclidean distance are usually non-convex. Solving these problems to obtain optimal p and q often requires complex calculations, and may even rely on singular value decomposition (SVD). This introduces a huge computational overhead in each iteration, reducing training efficiency.

[0182] To address the aforementioned problems in existing technologies, this invention provides a model training method based on alternating adaptive momentum optimization:

[0183] a) Reduce memory overhead to sublinear level: This invention aims to solve the problem of excessive memory consumption in the Adam algorithm by performing low-rank decomposition on the second moment, thereby reducing the optimizer's memory complexity from... Reduce to This allows for the training of larger-scale models with limited resources.

[0184] b) Introducing first-order moment estimation without additional memory overhead: To address the performance degradation of Adafactor due to the removal of momentum, this invention aims to propose a new mechanism that can effectively estimate and utilize first-order momentum without introducing additional memory overhead (i.e., without occupying additional mn GPU memory), thereby improving the convergence speed and robustness of the algorithm.

[0185] c) Employing an efficient alternating update strategy: Addressing the computational complexity of existing decomposition methods, this invention aims to utilize Euclidean distance as a more reasonable metric and employ an alternating minimization strategy to update the decomposition factors. This method avoids complex SVD calculations, ensuring that each update step reduces estimation errors while maintaining computational efficiency.

[0186] Specifically, this invention utilizes an alternating update rule based on Euclidean distance minimization to perform low-rank decomposition of the second-order momentum of the gradient. Unlike existing techniques (such as Adafactor) that use computationally complex SVD decomposition or physically ambiguous KL divergence, this invention updates the decomposition factors by alternately fixing row or column vectors in odd-even iteration steps, avoiding high computational costs. This strategy effectively reduces the storage requirements of the second-order momentum from the original... Significantly reduced to While ensuring that the decomposition error gradually decreases monotonically, it greatly improves the memory efficiency of large-scale matrix optimization, enabling it to support the training of larger-scale models.

[0187] Furthermore, addressing the issue of slow convergence and poor robustness in existing memory-efficient algorithms (such as Adafactor) due to the removal of first-order momentum, this invention designs a momentum preservation mechanism with zero additional memory. This mechanism utilizes the gradient accumulation characteristics of leaf nodes in automatic differentiation platforms (such as PyTorch) to directly reuse the space originally used for temporary gradient storage to persistently store first-order momentum, thus eliminating the need to allocate a separate GPU memory buffer. Simultaneously, this invention innovatively uses the bias-corrected squared first-order momentum to drive the update of second-order momentum, instead of using the original squared gradient. This allows the GPU memory occupied by the original gradient to be released immediately after the first-order momentum is calculated. This design enables the algorithm to achieve convergence speed and training robustness comparable to the standard Adam algorithm while maintaining extremely low memory usage.

[0188] To meet the needs of different deep learning model architectures, this invention provides a general tensor adaptation scheme. This scheme intelligently finds the optimal segmentation point, logically reshaping high-dimensional tensors of arbitrary dimensions (such as convolutional kernels in convolutional neural networks) into a two-dimensional matrix view with an aspect ratio as close as possible. Without incurring additional overhead from physical data copying, this scheme ensures that the aforementioned low-rank decomposition algorithm can be seamlessly applied to parameters of various shapes, and guarantees that the memory overhead from auxiliary variables remains negligible relative to the parameters themselves, thereby improving the algorithm's versatility and flexibility in complex model scenarios.

[0189] Compared with existing technologies, this invention (Alada) significantly reduces memory usage while maintaining excellent convergence speed and model generalization ability.

[0190] First, regarding memory efficiency, this invention employs a low-rank decomposition of second-order momentum using an alternating update strategy based on Euclidean distance minimization, and combines this with a first-order momentum mechanism that utilizes gradient attribute space reuse. This successfully reduces the optimizer's additional memory complexity from that of standard Adam. Reduce to This sublinear memory growth characteristic allows Alada to support deep learning models with larger training parameter scales under the same hardware conditions. Experimental data shows that in large language model training tasks such as GPT-2, the peak memory usage of this invention is reduced by more than 30% compared to Adam, and is on par with the extremely compressed Adafactor.

[0191] Secondly, regarding convergence speed and accuracy, this invention overcomes the training instability and slow convergence problems caused by Adafactor's discarding of first-order momentum. By introducing first-order momentum estimation with zero additional memory, Alada exhibits faster convergence speed and better generalization performance on natural language understanding (GLUE Benchmark) and machine translation tasks. Experimental results show that Alada outperforms Adafactor in BLEU scores and accuracy on most tasks, and can reach or even surpass the level of the standard Adam. Furthermore, theoretical analysis also proves that this invention possesses the same... The convergence rate ensures the robustness of the algorithm at the theoretical level.

[0192] In this embodiment of the invention, a model training method based on alternating adaptive momentum optimization is provided. The method involves obtaining initial model parameters and marking them as leaf nodes to construct a gradient space for storing first-order momentum. The initial model parameters are preprocessed based on their type, outputting target model parameters and analytical model parameter data. An unbiased first-order momentum calculation is performed using a predefined loss function based on a predefined training data distribution, target model parameters, and the first-order momentum of historical steps in the gradient space used to store first-order momentum. The unbiased first-order momentum of the current step is then output. Element-wise squaring is performed on the unbiased first-order momentum of the current step to output the second-order momentum of the current step. Finally, the alternating adaptive momentum method is used to optimize the model training based on the analytical model parameter data. The second-order momentum vector factor is updated alternately with odd and even values ​​for each step, determining the updated second-order momentum vector factor. An adaptive step size rule is used to iteratively update the target model parameters based on the model parameter parsing data, the updated second-order momentum vector factor, and the unbiased first-order momentum, obtaining the pre-trained optimal parameters. A pre-trained model is then constructed based on these optimal parameters. Based on this scheme, this invention constructs a dedicated gradient space to store the first-order momentum, reusing the gradient storage characteristics of leaf nodes to avoid additional memory usage. The preprocessing stage adapts to subsequent algorithms according to parameter type, reducing redundant storage overhead. Furthermore, this invention employs an alternating minimization strategy to update the second-order momentum vector factor alternately with odd and even values, eliminating the need to maintain a complete second-order momentum matrix of the same size as the parameter matrix, thus simplifying traditional adaptive optimization algorithms. The linear memory overhead is transformed into With its low-order overhead, it significantly improves the memory efficiency of large-scale matrix optimization and can support the training of larger-scale models.

[0193] Please see Figure 3 , Figure 3 This is a structural block diagram of a model training system based on alternating adaptive momentum optimization provided in Embodiment 2 of the present invention.

[0194] This invention provides a model training system based on alternating adaptive momentum optimization, comprising:

[0195] The acquisition module 301 is used to acquire the initial model parameters, mark the initial model parameters as leaf nodes, and construct a gradient space for storing the first-order momentum;

[0196] Preprocessing module 302 is used to preprocess the initial model parameters based on the model parameter type of the initial model parameters, and output the target model parameters and model parameter parsing data;

[0197] The calculation module 303 is used to perform unbiased first-order momentum calculation based on the predefined training data distribution, target model parameters and the first-order momentum of the historical steps in the gradient space used to store the first-order momentum using a predefined loss function, and output the unbiased first-order momentum of the current step.

[0198] The second-order momentum output module 304 is used to perform element-wise squaring of the unbiased first-order momentum of the current step and output the second-order momentum of the current step.

[0199] The update module 305 is used to update the second-order momentum vector factor of the current step by alternating odd and even according to the model parameter analytical data using the alternating adaptive momentum method, and to determine the updated second-order momentum vector factor.

[0200] Module 306 is used to iteratively update the target model parameters using an adaptive step size rule based on the model parameter parsing data, the updated second-order momentum vector factor, and the unbiased first-order momentum, to obtain the pre-trained optimal parameters, and to construct a pre-trained model based on the pre-trained optimal parameters.

[0201] Those skilled in the art will understand that, for the sake of convenience and brevity, the specific working process of the system and modules described above can be referred to the corresponding process in the foregoing method embodiments, and will not be repeated here.

[0202] This invention also provides a computer device, including a memory and a processor, wherein the memory stores a computer program; when the computer program is executed by the processor, the processor performs the steps of the model training method based on alternating adaptive momentum optimization as described in the above embodiments.

[0203] This invention also provides a computer-readable storage medium storing a computer program / instructions thereon, which, when executed by a processor, implements the steps of the model training method based on alternating adaptive momentum optimization as described in the above embodiments.

[0204] This invention also provides a computer program product, including a computer program stored on a non-transitory computer-readable storage medium, the computer program including program instructions, wherein when the program instructions are executed by a computer, the computer performs the steps of the model training method based on alternating adaptive momentum optimization as described in the above embodiments.

[0205] In the several embodiments provided in this application, it should be understood that the disclosed systems and methods can be implemented in other ways. For example, the device embodiments described above are merely illustrative; for instance, the division of units is only a logical functional division, and in actual implementation, there may be other division methods. For example, multiple units or components may be combined or integrated into another system, or some features may be ignored or not executed. Furthermore, the coupling or direct coupling or communication connection shown or discussed may be through some interfaces; the indirect coupling or communication connection between devices or units may be electrical, mechanical, or other forms.

[0206] The units described as separate components may or may not be physically separate. The components shown as units may or may not be physical units; that is, they may be located in one place or distributed across multiple network units. Some or all of the units can be selected to achieve the purpose of this embodiment according to actual needs.

[0207] Furthermore, the functional units in the various embodiments of the present invention can be integrated into one processing unit, or each unit can exist physically separately, or two or more units can be integrated into one unit. The integrated unit can be implemented in hardware or as a software functional unit.

[0208] If the integrated unit is implemented as a software functional unit and sold or used as an independent product, it can be stored in a computer-readable storage medium. Based on this understanding, the technical solution of the present invention, in essence, or the part that contributes to the prior art, or all or part of the technical solution, can be embodied in the form of a software product. This computer software product is stored in a storage medium and includes several instructions to cause a computer device (which may be a personal computer, server, or network device, etc.) to execute all or part of the steps of the methods of the various embodiments of the present invention. The aforementioned storage medium includes various media capable of storing program code, such as USB flash drives, portable hard drives, read-only memory (ROM), random access memory (RAM), magnetic disks, or optical disks.

[0209] The above embodiments are only used to illustrate the technical solutions of the present invention, and are not intended to limit it. Although the present invention has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that modifications can still be made to the technical solutions described in the foregoing embodiments, or equivalent substitutions can be made to some of the technical features. Such modifications or substitutions do not cause the essence of the corresponding technical solutions to deviate from the spirit and scope of the technical solutions of the embodiments of the present invention.

Claims

1. A model training method based on alternating adaptive momentum optimization, characterized in that, The method comprises the following steps: acquiring initial model parameters and marking the initial model parameters as leaf nodes to build a gradient space for storing first-order momentum; preprocessing the initial model parameters based on the model parameter type of the initial model parameters to output target model parameters and model parameter analysis data; performing unbiased first-order momentum calculation on the target model parameters and the first-order momentum of the historical steps in the gradient space for storing first-order momentum according to a predefined loss function and a predefined training data distribution to output the unbiased first-order momentum of the current step; performing element-by-element squaring operation on the unbiased first-order momentum of the current step to output the second-order momentum of the current step; performing odd-even alternating update on the vector factor of the second-order momentum of the current step according to the model parameter analysis data by using an alternating adaptive momentum method to determine the updated second-order momentum vector factor; performing iterative update on the target model parameters according to the model parameter analysis data, the updated second-order momentum vector factor and the unbiased first-order momentum by using an adaptive step length rule to obtain pre-training optimal parameters, and constructing a pre-training model according to the pre-training optimal parameters.

2. The model training method based on alternating adaptive momentum optimization according to claim 1, characterized in that, The preprocessing of the initial model parameters based on the model parameter type of the initial model parameters to output target model parameters and model parameter analysis data comprises the following steps: if the model parameter type of the initial model parameters is in the form of a matrix, recording the number of rows and columns of the initial model parameters to obtain model parameter analysis data, and taking the initial model parameters as target model parameters; if the model parameter type of the initial model parameters is a high-dimensional tensor, recording the dimension size of the initial model parameters to obtain model parameter analysis data, and performing view reshaping on the initial model parameters to obtain target model parameters. 3.The model training method based on alternating adaptive momentum optimization of claim 1, wherein, The unbiased first-order momentum calculation on the target model parameters and the first-order momentum of the historical steps in the gradient space for storing first-order momentum according to a predefined loss function and a predefined training data distribution to output the unbiased first-order momentum of the current step comprises the following steps: randomly extracting a training sample in the predefined training data distribution, taking the training sample as the input of a model corresponding to the target model parameters to obtain an output result; calculating the loss value of the current step based on the predefined loss function according to the output result; calculating the random gradient of the loss value of the current step with respect to the target model parameters based on the loss value of the current step; performing in-place scaling on the first-order momentum of the historical steps in the gradient space for storing first-order momentum and the random gradient to obtain the scaled first-order momentum of the historical steps and the scaled random gradient; accumulating the scaled random gradient into the scaled first-order momentum of the historical steps to obtain the first-order momentum of the current step; performing bias correction on the first-order momentum of the current step to obtain the unbiased first-order momentum of the current step.

4. The model training method based on alternating adaptive momentum optimization according to claim 1, characterized in that, The vector factor includes a row vector factor and a column vector factor; the updated second-order momentum vector factor includes a row vector factor of a current step and a column vector factor of the current step; and the adopting the alternating adaptive momentum method to update the vector factor of the second-order momentum of the current step according to the model parameter analysis data to determine the updated second-order momentum vector factor includes: determining whether the current step is an initial iteration step; if yes, calculating an initial gradient norm mean based on dimension information of the model parameter analysis data; setting each element value of the row vector factor and the column vector factor of the second-order momentum of the current step to a square root of the initial gradient norm mean to obtain an initial row vector factor and an initial column vector factor; taking the initial row vector factor and the initial column vector factor as the row vector factor of the corresponding current step and the column vector factor of the current step; if no, determining whether the current step is odd or even; if the current step is even, fixing the column vector factor unchanged and taking the column vector factor as the column vector factor of the current step; performing matrix and vector multiplication, norm calculation and division on the row vector factor in sequence to obtain a row vector factor optimal analytical solution; performing exponential moving average on the row vector factor optimal analytical solution and the updated row vector factor of a historical step to obtain the row vector factor of the current step; if the current step is odd, fixing the row vector factor unchanged and taking the row vector factor as the row vector factor of the current step; performing matrix transposition, matrix and vector multiplication, norm calculation and division on the column vector factor in sequence to obtain a column vector factor optimal analytical solution; performing exponential moving average on the column vector factor optimal analytical solution and the updated column vector factor of the historical step to obtain the column vector factor of the current step.

5. The model training method based on alternating adaptive momentum optimization according to claim 1, characterized in that, The adopting the adaptive step length rule to update the target model parameter according to the model parameter analysis data, the updated second-order momentum vector factor and the unbiased first-order momentum to obtain a pre-training optimal parameter includes: reconstructing a second-order moment matrix based on the updated second-order momentum vector factor; performing magnitude correction and bias correction on the second-order moment matrix in sequence to obtain a corrected second-order moment; adding a numerical stability constant to the corrected second-order moment to obtain a corrected second-order moment with a stability term; updating the target model parameter according to the adaptive step length rule based on the unbiased first-order momentum and the corrected second-order moment with the stability term to obtain an updated target model parameter; when the iteration termination condition is reached, if a model parameter type of the updated target model parameter determined when the iteration termination condition is reached is a high-dimensional tensor, restoring the updated target model parameter determined when the iteration termination condition is reached to an original tensor shape according to the model parameter analysis data, verifying the dimensions to be correct and determining the pre-training optimal parameter; if the model parameter type of the updated target model parameter determined when the iteration termination condition is reached is a matrix form, taking the updated target model parameter determined when the iteration termination condition is reached as the pre-training optimal parameter.

6. The model training method based on alternating adaptive momentum optimization according to claim 5, characterized in that, Further comprising: When the iteration termination condition is not reached, the updated target model parameter is taken as a new target model parameter, and the step of performing the unbiased first-order momentum calculation according to the pre-defined training data distribution, the target model parameter and the first-order momentum of the historical step in the gradient space for storing the first-order momentum is jumped to execute, and the unbiased first-order momentum of the current step is output until the iteration termination condition is reached.

7. A model training system based on alternating adaptive momentum optimization, the system comprising: The method comprises the steps of: an acquisition module, configured to acquire an initial model parameter, and mark the initial model parameter as a leaf node, and construct a gradient space for storing a first-order momentum; a preprocessing module, configured to preprocess the initial model parameter based on a model parameter type of the initial model parameter, and output a target model parameter and model parameter analysis data; a calculation module, configured to perform an unbiased first-order momentum calculation according to a pre-defined training data distribution, the target model parameter and the first-order momentum of a historical step in the gradient space for storing the first-order momentum by using a pre-defined loss function, and output an unbiased first-order momentum of a current step; a second-order momentum output module, configured to perform an element-by-element squaring operation on the unbiased first-order momentum of the current step, and output a second-order momentum of the current step; an update module, configured to perform an odd-even alternating update on a vector factor of the second-order momentum of the current step according to the model parameter analysis data by using an alternating adaptive momentum method, and determine an updated second-order momentum vector factor; a construction module, configured to perform an iterative update on the target model parameter according to the model parameter analysis data, the updated second-order momentum vector factor and the unbiased first-order momentum by using an adaptive step rule, obtain a pre-training optimal parameter, and construct a pre-training model according to the pre-training optimal parameter.

8. An electronic device, comprising: The computer program is executed to implement the model training method based on the alternating adaptive momentum optimization.

9. A computer readable storage medium having stored thereon a computer program, characterized in that, The computer program is executed to implement the model training method based on the alternating adaptive momentum optimization.

10. A computer program product, characterised in that, The computer program product comprises a computer program stored on a non-transitory computer-readable storage medium, and the computer program comprises program instructions, wherein when the program instructions are executed by a computer, the computer is caused to execute the steps of the model training method based on the alternating adaptive momentum optimization. The computer program product comprises a computer program stored on a non-transitory computer-readable storage medium, and the computer program comprises program instructions, wherein when the program instructions are executed by a computer, the computer is caused to execute the steps of the model training method based on the alternating adaptive momentum optimization.

Images