A vector quantization diffusion method for mechanical fault diagnosis data enhancement

Abstract

The application belongs to the technical field of intelligent fault diagnosis and data enhancement of rotating machinery, and relates to a vector quantization diffusion method for mechanical fault diagnosis data enhancement, wherein a frequency domain representation input vector of a training sample in a mechanical fault diagnosis training set is input into an encoder in a vector quantization module to output a continuous latent variable; the continuous latent variable is input into a codebook quantization unit in the vector quantization module to output a quantized latent representation of the continuous latent variable; the quantized latent representation is input into a decoder in the vector quantization module to output a reconstructed sample; a reconstruction loss, a codebook embedding loss, a commitment loss and a codebook utilization rate regularization term are constructed; a quantization loss function is constructed based on the reconstruction loss, the codebook embedding loss, the commitment loss and the codebook utilization rate regularization term, so that the vector quantization module is iteratively trained, and the trained vector quantization module is used to reconstruct the training sample.

Description

Technical Field

[0001] This invention relates to the field of intelligent fault diagnosis and data augmentation technology for rotating machinery, and in particular to a vector quantization diffusion method for data augmentation in mechanical fault diagnosis. Background Technology

[0002] With the rapid development of deep learning, researchers have successfully applied Intelligent Fault Diagnosis (IFD) technology to fault identification tasks in industrial equipment such as rotating machinery. Deep learning-based IFD models, with their powerful feature learning capabilities, have become a widely used and efficient method. Although IFD technology has made significant progress, its successful application relies heavily on large-scale and evenly distributed datasets. In actual industrial production, the normal operating time of mechanical equipment far exceeds the time of failure, resulting in a much smaller scale of fault samples collected by sensors compared to normal samples. Simultaneously, manually labeling large amounts of data is not only time-consuming but also incurs high labor costs, severely limiting the efficiency and scale of labeled data acquisition. This imbalance in sample distribution seriously hinders the further application of IFD technology.

[0003] To address these challenges, data augmentation techniques, which generate synthetic samples to augment minority class data, have become an effective solution. Currently, methods for addressing the imbalanced sample problem are mainly based on two categories: Generative Adversarial Networks (GANs) and diffusion models. GAN methods include DCGAN, CGAN, and WGAN, while diffusion model methods include DDPM and its improved variants. These methods balance the distribution of samples from different health states by generating missing samples, providing more balanced training data for diagnostic models and demonstrating effectiveness in improving diagnostic accuracy and generalization performance. However, existing methods still have limitations: GAN-type methods, due to their adversarial generative nature, are prone to stability issues such as gradient vanishing and pattern collapse during training. At the same time, the samples generated by GAN-type models are highly similar to the training data, resulting in a lack of diversity in fault features and difficulty in effectively expanding rare fault samples. Diffusion model-type methods suffer from slow sampling speed, making it impossible to efficiently generate a large number of fault samples. In addition, the noise prediction models of diffusion models are mostly built on the U-Net architecture, which has limited feature extraction capabilities and cannot fully explore the unique time-frequency distribution patterns in fault signals. This results in a large deviation between the generated fault samples and the time-frequency features of real industrial fault signals, making it difficult to support the training and accurate identification of subsequent fault diagnosis models.

[0004] In summary, existing methods for data augmentation in mechanical fault diagnosis suffer from the problem of failing to balance the quantity and quality of generated samples. They cannot generate a sufficient number of highly diverse missing samples to construct a balanced dataset, resulting in insufficient training data for subsequent fault diagnosis models and affecting diagnostic accuracy. Summary of the Invention

[0005] Therefore, the technical problem to be solved by the present invention is to overcome the problem that existing methods for data augmentation in mechanical fault diagnosis have difficulty in balancing the quantity and quality of sample generation, and cannot generate a sufficient number of missing samples with high diversity to build a balanced dataset, resulting in insufficient training data for subsequent fault diagnosis models and affecting diagnostic accuracy.

[0006] To address the aforementioned technical problems, this invention provides a vector quantization diffusion method for enhancing mechanical fault diagnosis data, comprising: Obtain training samples from the mechanical fault diagnosis training set and convert the training samples into a frequency domain representation; The frequency domain representation is encoded by the encoder in the vector quantization module, and the output is a continuous latent variable. The continuous latent variable is input into the codebook quantization unit in the vector quantization module. The Euclidean distance between the continuous latent variable and all codewords in the learnable codebook is calculated. The quantized latent representation of the continuous latent variable is obtained based on the codeword index with the smallest Euclidean distance. The quantized latent representation is input into the decoder in the vector quantization module for decoding, and the reconstructed sample is output. The reconstruction loss is constructed based on the errors of training samples and reconstructed samples; the codebook embedding loss is constructed based on the errors of the stopping gradient terms of quantized latent representation and continuous latent variables; the commitment loss is constructed based on the errors of the stopping gradient terms of continuous latent variables and quantized latent representation; and the codebook utilization regularization term is constructed based on the usage frequency of each codeword in the learnable codebook. A quantization loss function is constructed based on reconstruction loss, codebook embedding loss, commitment loss, and codebook utilization regularization term, thereby iteratively training the vector quantization module and using the trained vector quantization module to reconstruct the training samples.

[0007] Preferably, when iteratively training the vector quantization module, a pass-through estimator is used to perform quantization mapping during forward propagation and to construct an encoder input approximate gradient channel during backpropagation, thereby using the pass-through approximation method to transmit gradients and perform backpropagation gradient updates.

[0008] Preferably, continuous latent variables Represented as: , in, , , , and These represent the height and width of the frequency domain representation, respectively. and This represents the height and width of the frequency domain representation after downsampling. Indicates the number of downsampling attempts; The encoder consists of multiple convolutional layers and residual blocks, and compresses the spatial dimension through downsampling operations. The frequency domain representation of the training samples; Quantitative latent representation for: , , in, Represents the position in the quantized latent representation The quantization value at that location; This represents the k-th codeword in the learnable codebook; This indicates the number of codewords in the learnable codebook; Represents position in continuous latent variables The value of the variable at that location; This represents the codeword index with the smallest Euclidean distance.

[0009] Preferably, the quantization loss function Represented as: , , , , , in, Indicates the reconstruction loss; Indicates training samples; Indicates the reconstructed sample; Indicates codebook embedding loss; Represents the stopping gradient term of a continuous latent variable. This indicates the stopping gradient operator during forward propagation. During back propagation ; express The weights; Indicates a commitment to loss; This represents the stopping gradient term of the quantized latent representation; express The weights; Regular expressions representing codebook utilization; This represents the frequency of use of the k-th codeword in the learnable codebook; express The weight.

[0010] Preferably, after iterative training of the vector quantization module, the method further includes: The training samples are input into the encoder in the trained vector quantization module for encoding, and the output is a continuous latent variable. The continuous latent variable is input into the codebook quantization unit in the trained vector quantization module. The Euclidean distance between the continuous latent variable and all codewords in the learnable codebook is calculated. The quantized latent representation of the continuous latent variable is obtained based on the codeword index with the smallest Euclidean distance. The quantized latent representation is input into the noise addition unit in the diffusion generation module to add noise, resulting in a noisy representation. The Transformer encoder in the noisy representation input diffusion generation module is modeled to convert the noisy representation into a noisy sequence. The noise-added sequence, the true fault category of the training samples, and the time step corresponding to the diffusion step size are input into the noise prediction unit in the diffusion generation module for noise prediction and denoising. The noise prediction error is used to construct a noise prediction loss function, thereby iteratively training the diffusion generation module to obtain a trained diffusion generation module for denoising the quantized latent representation and outputting the target quantized latent representation.

[0011] Preferably, noise is added. for: , in, This represents the cumulative decay coefficient from the initial time to step t; This represents real noise and conforms to a Gaussian distribution. This represents the quantized latent representation output by the codebook quantization unit in the trained vector quantization module; Noise prediction loss function Represented as: , in, This represents the expectation of a fully random sample. Indicates predicted noise; Indicates a time step; Indicates the actual fault category.

[0012] Preferably, the noise-added sequence, the true fault category of the training samples, and the time step corresponding to the diffusion step size are input into the noise prediction unit in the diffusion generation module for noise prediction and denoising, including: The noisy sequence is sequentially subjected to patch segmentation and linear transformation, feature recombination and linear transformation to obtain serialized input features; The true fault categories of the training samples are input into the label embedding layer and mapped to obtain the category conditional embedding vector. The diffusion time step is then processed sequentially through temporal embedding, linear transformation, ReLU activation, and quadratic linear mapping to generate temporal features. The category conditional embedding vector and the temporal features are fused to obtain the conditional embedding features. The serialized input features and conditional embedding features are input into two concatenated diffusion transform blocks to obtain global time-frequency features; The global time-frequency features are sequentially processed through layer normalization, linear dimension transformation, and feature recombination to obtain the noise corresponding to the current time step for single-step noise estimation. The noise is then gradually removed by iterative multi-time-step reverse diffusion.

[0013] Preferably, the serialized input features and conditional embedding features are input into a two-layer concatenated diffusion transform block to obtain global time-frequency features, including: The serialized input features are used as the input features of the first layer diffusion transform block, and the first adaptive layer normalization is used to normalize the input features to obtain the first normalized features. The first normalized feature is scaled and shifted and modulated using the linearly transformed conditional embedding feature to obtain the first modulated feature; The original time-domain signal of the training samples is mapped to a time-domain feature representation with the same dimension as the quantized latent representation. The time-domain feature representation and the first modulation feature are input into the cross-attention layer, and the output is the time-series correlation feature. The scaled temporal correlation features and input features are adaptively normalized using the second adaptive layer normalization to obtain the second normalized features; The second normalized feature is scaled and shifted to obtain the second modulation feature; The second modulation feature is input into the feedforward fully connected layer for feature mapping. The mapped feature output by the feedforward fully connected layer is scaled. The output of the first layer diffusion transform block is obtained based on the scaled mapped feature and the scaled temporal correlation feature. The output of the first-layer diffusion transform block is used as the input feature of the second-layer diffusion transform block, and the global time-frequency features are obtained based on the output of the second-layer diffusion transform block.

[0014] Preferably, the temporal feature representation and the first modulation feature are input into the cross-attention layer, and the output temporal correlation features include: The first modulation feature is used as the query vector, and the temporal feature representation is used as the key vector and value vector to calculate the cross-attention weights. By using cross-attention weights to weight and aggregate the temporal feature representations, temporal correlation features are obtained.

[0015] Preferably, after obtaining the trained diffusion generation module, the method further includes: The target quantization latent representation output by the trained diffusion generation module is input into the codebook quantization unit in the trained vector quantization module. The Euclidean distance between the target quantization latent representation and all codewords in the learnable codebook is calculated. The quantization latent variable of the target quantization latent representation is obtained based on the codeword index with the smallest Euclidean distance. The latent variables are quantized and input into the decoder in the trained vector quantization module for decoding, and the output is diffused to generate samples.

[0016] The vector quantization diffusion method for enhancing mechanical fault diagnosis data provided in this application has the following beneficial effects: 1. A vector quantization autoencoder structure adapted to rotating machinery fault signals was designed. This structure consists of an encoder with residual skip connections, a learnable codebook, and a decoder. Deep features of time-frequency representation are extracted through downsampling convolution and residual modules, compressed into a continuous latent space representation, and mapped from the continuous latent space to the discrete latent space using a nearest neighbor search strategy. This fully preserves the key impact and modulation features of the fault signal, accurately capturing the core feature patterns of different fault types. Compared to existing enhancement methods that only make minor perturbations on the original samples, this application generates novel missing fault samples with near-real distribution through codeword combination, achieving sufficient and highly diverse data enhancement. Furthermore, a multi-loss joint optimization strategy for the vector quantization module is proposed. This strategy ensures the quality of sample reconstruction through reconstruction loss, achieves adaptive learning of the codebook vector through embedding loss, and constrains the fluctuation amplitude of the encoder output through commitment loss. The weighted joint optimization of these three losses improves the stability of latent space modeling and the ability to express fault features, thereby generating high-quality fault samples. 2. Building upon vector quantization, this application introduces a diffusion generation mechanism and designs a diffusion model noise prediction architecture based on DiT Blocks. This architecture employs a two-layer cascaded DiT Block structure with an adaptive normalization layer. Scaling and translation factors are dynamically generated using a lightweight multilayer perceptron, deeply integrating time step embedding and fault category label conditional information into the noise prediction process. This achieves targeted generation of specific fault categories, overcoming the parameter constraints of traditional static normalization and the global feature extraction bottleneck of the U-Net architecture. Simultaneously, a time-frequency domain cross-attention mechanism for fault signals is designed. This mechanism uses the frequency domain latent variable token as the query and the time-domain features extracted from the original one-dimensional vibration signal through a lightweight convolutional coding network as the key and value. This achieves explicit interaction and deep fusion of time-domain dynamic characteristics and frequency-domain structural information, enhancing the model's ability to mine time-frequency domain correlation features of fault signals and improving noise prediction accuracy and the realism of generated samples. In addition, a diffusion preprocessing and conditional embedding method adapted to discrete latent variables is proposed. This method divides the discrete latent variables output by vector quantization into patch sequences and completes linear mapping. It combines learnable position encoding, time step embedding and category conditional embedding to adapt to the sequence input requirements of Transformer. At the same time, it adopts the Cosine noise scheduling strategy to complete the noise processing of latent variables, ensuring the stability and controllability of the diffusion generation process. 3. This application designs a three-level data augmentation architecture for rotating machinery fault diagnosis, consisting of time-frequency representation, vector quantization, and diffusion generation. It converts one-dimensional vibration time-series signals into two-dimensional time-frequency representations using short-time Fourier transform. After mapping to a discrete latent space via a vector quantization module, it completes latent space diffusion generation and decoding reconstruction based on a Diffusion Transformer architecture with time-frequency cross-attention, achieving targeted, high-quality generation of minority class fault samples and completing dataset equalization. Simultaneously, a two-stage model training and sample generation method is employed. First, a vector quantization autoencoder is pre-trained using source domain fault data to complete the latent space mapping modeling of the fault signal. Then, with fixed vector quantization module parameters, a DiT-based noise prediction network is trained. Finally, discrete latent variables are generated through iterative denoising, and the decoder reconstructs the fault samples. This solves the problem of mutual interference between latent space mapping and diffusion generation processes during end-to-end training, ensuring the stability of model training. Attached Figure Description

[0017] To make the content of this invention easier to understand, the invention will be further described in detail below with reference to specific embodiments and accompanying drawings, wherein: Figure 1 A flowchart of the vector quantization diffusion method for enhancing mechanical fault diagnosis data provided in this application; Figure 2 A schematic diagram illustrating the structure and optimization principle of the vector quantization module provided in this application; Figure 3 A schematic diagram of the noise prediction unit and diffusion conversion block provided in this application; Figure 4 A schematic diagram illustrating the generation process of the diffusion-generated sample provided in this application; Figure 5 A schematic diagram illustrating the fault diagnosis accuracy of different sample generation methods provided in this application. Detailed Implementation

[0018] The present invention will be further described below with reference to the accompanying drawings and specific embodiments, so that those skilled in the art can better understand and implement the present invention. However, the embodiments described are not intended to limit the present invention.

[0019] Please refer to the following: Figure 1 and Figure 2 , Figure 1 The diagram shown is a flowchart of the vector quantization diffusion method for enhancing mechanical fault diagnosis data provided in this application. Figure 2 The diagram shown illustrates the structure and optimization principle of the vector quantization module provided in this application. The method specifically includes steps S10 to S60: S10: Obtain training samples from the mechanical fault diagnosis training set and convert the training samples into frequency domain representation.

[0020] For example, the training samples are the original time series signals and their true fault category labels. After the original time series signals are truncated by a sliding window, a short-time Fourier transform is performed. Then, through operations such as regularization and geometric transformation, the one-dimensional time series data is converted into a two-dimensional time-frequency representation, thereby obtaining the time-domain and frequency-domain representation of the training samples.

[0021] S20: Encode the frequency domain representation input vector quantization module into the encoder to output a continuous latent variable.

[0022] Specifically, continuous latent variables Represented as: , in, , , , and These represent the height and width of the frequency domain representation, respectively. and This represents the height and width of the frequency domain representation after downsampling. Indicates the number of downsampling attempts; The encoder consists of multiple convolutional layers and residual blocks, and compresses the spatial dimension through downsampling operations. The frequency domain representation of the training samples. Indicates the number of channels.

[0023] For example, encoder By compressing the spatial dimension through two downsampling operations with a step size of 2, the frequency domain representation with an input size of 32*32 is mapped to an 8*8 latent feature representation, thereby significantly reducing the computational complexity of the subsequent calculations.

[0024] S30: Input the continuous latent variable into the codebook quantization unit in the vector quantization module, calculate the Euclidean distance between the continuous latent variable and all codewords in the learnable codebook, and obtain the quantized latent representation of the continuous latent variable based on the codeword index with the smallest Euclidean distance.

[0025] Specifically, quantizing latent representations for: , , in, Represents the position in the quantized latent representation The quantization value at that location; This represents the k-th codeword in the learnable codebook; This indicates the number of codewords in the learnable codebook; Represents position in continuous latent variables The value of the variable at that location; This represents the codeword index with the smallest Euclidean distance.

[0026] Specifically, in order to realize the mapping from continuous latent space to discrete latent space, this application introduces a learnable codebook of size K. Each codeword is a vector of dimension D. This learnable codebook maps continuous latent variables to a finite number of discrete vectors, thus achieving a discretized representation of the latent space. For each spatial location of a continuous latent variable... By using a nearest neighbor search strategy for vector quantization, continuous latent scalars are ultimately mapped to a finite codebook space, resulting in a clearer structured semantic representation and enhancing the stability and reliability of the generation process.

[0027] S40: Decode the latent representation input to the decoder in the vector quantization module and output the reconstructed sample.

[0028] S50: Construct reconstruction loss based on the errors of training samples and reconstructed samples; construct codebook embedding loss based on the errors of quantized latent representation and stopping gradient terms of continuous latent variables; construct commitment loss based on the errors of stopping gradient terms of continuous latent variables and quantized latent representation; and construct codebook utilization regularization term based on the usage frequency of each codeword in the learnable codebook.

[0029] S60: Construct a quantization loss function based on reconstruction loss, codebook embedding loss, commitment loss, and codebook utilization regularization term, thereby iteratively training the vector quantization module and using the trained vector quantization module to reconstruct the training samples.

[0030] Specifically, the quantization loss function Represented as: , , , , , in, Indicates the reconstruction loss; Indicates training samples; Indicates the reconstructed sample; Indicates codebook embedding loss; Represents the stopping gradient term of a continuous latent variable. This indicates the stopping gradient operator during forward propagation. During back propagation ; express The weights; Indicates a commitment to loss; This represents the stopping gradient term of the quantized latent representation; express The weights; Regular expressions representing codebook utilization; This represents the frequency of use of the k-th codeword in the learnable codebook; express The weight.

[0031] Reconstruction loss measures the pixel-level difference between the decoder output and the original input, driving the decoder to learn its ability to reconstruct high-quality samples from discrete latent variables. Codebook embedding loss constrains the quantized latent representation to align with the encoder output, enabling adaptive learning of the vector quantization dictionary. This loss freezes the gradients of the continuous latent variables in the encoder output, allowing only the codebook quantization unit to receive gradient updates, pushing the selected codeword closer to its corresponding continuous latent variable. This ensures the codebook covers the true distribution of the training data while preventing interference with the encoder parameters. Commitment loss limits the fluctuation range of the encoder output, forcing it to commit to the currently selected codeword. This loss is achieved by freezing the quantization... The gradient of the latent representation is updated only by the encoder to prevent frequent jumps in the encoder output, enhance training stability, and avoid oscillations in the optimization process caused by drastic changes in the codeword index. On this basis, in order to further enhance the structured representation capability of the quantized latent representation and optimize the efficiency of latent space modeling, this application also introduces an efficient symbol distribution regulation mechanism. That is, the statistical distribution of the quantization index is explicitly modeled using a codebook utilization regularization term, and an information entropy-based constraint term is incorporated to optimize the codebook activation mode and suppress the codebook collapse phenomenon. This regularization term enhances the structural richness of the quantized latent representation by promoting the balanced distribution of discrete symbols and provides a stable discrete prior foundation.

[0032] Optionally, to ensure end-to-end training, this application also introduces a straight-through estimator (STE) to implement approximate gradient backpropagation. Specifically, during iterative training of the vector quantization module, the straight-through estimator is used to perform quantization mapping during forward propagation and to construct an approximate gradient channel for the encoder input during backpropagation. This allows for gradient propagation using a straight-through approximation method for backpropagation gradient updates, as expressed in the formula: ,in, This represents the pass-through equivalent output characteristic, which passes the approximate gradient to the encoder output during backpropagation.

[0033] Furthermore, after completing the vector quantization modeling of the frequency domain representation, this application introduces a diffusion generation mechanism in the discrete latent variable space to achieve high-quality synthesis of minority class fault samples.

[0034] Specifically, after iterative training of the vector quantization module, steps 1 through 5 are also included: Step 1: Input the training samples into the encoder in the trained vector quantization module for encoding, and output continuous latent variables.

[0035] Step 2: Input the continuous latent variable into the codebook quantization unit in the trained vector quantization module, calculate the Euclidean distance between the continuous latent variable and all codewords in the learnable codebook, and obtain the quantized latent representation of the continuous latent variable based on the codeword index with the smallest Euclidean distance.

[0036] Step 3: Input the quantized latent representation into the noise addition unit in the diffusion generation module to add noise, and obtain the noisy representation.

[0037] Specifically, the quantized latent representation is used as input to the diffusion generation module, and noise is added progressively using a Cosine noise scheduling strategy. At time step t, the standard deviation of the added noise is... If the diffusion step number T = 1000, then the noise addition at step t is expressed as: , in, This represents the cumulative decay coefficient from the initial time to step t; Represents real noise, conforms to Gaussian distribution; This represents the quantized latent representation output by the codebook quantization unit in the trained vector quantization module.

[0038] Step 4: Model the Transformer encoder in the noisy representation input diffusion generation module to convert the noisy representation into a noisy sequence.

[0039] Specifically, this application employs a Transformer encoder structure to convert the noisy representation into a sequence form to adapt to the subsequent noise prediction unit. Formally, the noisy representation is divided into segments in the corresponding channels through convolution. Each size is The spatial patch is then flattened into a one-dimensional vector and mapped to a fixed-dimensional embedding space through a linear projection layer. Specifically, the patch size is set to 8, which is the input... The data is processed into four sizes using a convolution operation. The patch, flatten it out as The model uses tokens, each containing partial features of the input signal. After patch embedding, a learnable positional encoding is introduced to record spatial location information, ensuring the model simultaneously perceives local features and location context. Simultaneously, the time step embedding `t` is injected as conditional information into the encoder network, guiding the model to model the stage-specific features during the denoising process. Leveraging a multi-layer Transformer encoder structure, the model can capture global dependencies between different spatial locations in the latent space, thereby enhancing the structural consistency and spectral distribution rationality of the generated samples.

[0040] Step 5: Input the noise-added sequence, the true fault category of the training samples, and the time step corresponding to the diffusion step size into the noise prediction unit in the diffusion generation module for noise prediction and denoising. Construct a noise prediction loss function with the noise prediction error, and then iteratively train the diffusion generation module to obtain a trained diffusion generation module to denoise the quantized latent representation and output the target quantized latent representation.

[0041] Specifically, to achieve the generation of minority class samples for a specified fault category, this application introduces a category conditional embedding mechanism in the diffusion generation module. This mechanism maps the true fault category label to a learnable embedding vector, which, along with the time-step embedding, is fed into the noise prediction unit as conditional information. This guides the diffusion generation module to approximate the potential distribution of the fault category during the reverse denoising stage. Therefore, the noise prediction loss function... Represented as: , in, This represents the expectation of a fully random sample. Indicates predicted noise; Indicates a time step; Indicates the actual fault category.

[0042] By introducing diffusion generation on top of vector quantization, the stability and quality of minority class sample generation are enhanced while significantly reducing computational complexity.

[0043] Noise prediction is crucial to the quality of diffusion-generated samples. This application designs a Diffusion Transformer architecture containing two layers of DiT Blocks to accomplish the noise prediction task, such as... Figure 3 The diagram shown is a structural schematic of the noise prediction unit and diffusion transformation block provided in this application.

[0044] Specifically, the noise-added sequence, the true fault category of the training samples, and the time step corresponding to the diffusion step size are input into the noise prediction unit in the diffusion generation module for noise prediction and denoising, including steps 1-1 to 1-4: Step 1-1: After performing Patch segmentation and linear transformation, feature recombination and linear transformation on the noisy sequence, the serialized input features are obtained.

[0045] Steps 1-2: Input the true fault categories of the training samples into the label embedding layer and map them to obtain the category conditional embedding vector. After the diffusion time step is processed by temporal embedding, linear transformation, ReLU activation and quadratic linear mapping in sequence, the temporal features are generated. The category conditional embedding vector and the temporal features are fused to obtain the conditional embedding features.

[0046] Steps 1-3: Input the serialized input features and the conditional embedding features into a two-layer concatenated diffusion transform block to obtain global time-frequency features.

[0047] Steps 1-4: After the global time-frequency features are successively subjected to layer normalization, linear dimension transformation and feature recombination, the noise corresponding to the current time step is obtained for single-step noise estimation. The noise is then gradually removed by iterative multi-time-step reverse diffusion.

[0048] Furthermore, such as Figure 3 As shown, the global time-frequency features are obtained by cascading two layers of diffusion transform blocks, namely the serialized input features and the conditional embedding features, including steps 2-1 to 2-7: Step 2-1: Use the serialized input features as the input features of the first layer diffusion transform block, and use the first adaptive layer normalization to normalize the input features to obtain the first normalized features.

[0049] Step 2-2: Scale and shift the first normalized feature using the linearly transformed conditional embedding feature to obtain the first modulation feature.

[0050] Steps 2-3: Map the original time-domain signal of the training samples to a time-domain feature representation with the same dimension as the quantized latent representation. Input the time-domain feature representation and the first modulation feature into the cross-attention layer and output the temporal correlation feature.

[0051] Specifically, the original time-domain signal First, it is mapped via a lightweight encoding network to a temporal feature representation that matches the dimension of the quantized latent representation token. : , in, This represents a lightweight convolutional coding network that uses layer-by-layer downsampling and nonlinear transformation to compress temporal signals into feature representations that match the dimension of the quantized latent representation token, thus supporting subsequent cross-attention modeling. This design effectively preserves the key dynamic characteristics of temporal signals while constraining computational overhead.

[0052] Steps 2-4: Use the second adaptive layer normalization pair to adaptively normalize the scaled temporal correlation features and input features to obtain the second normalized features.

[0053] Steps 2-5: Scale and translate the second normalized feature to obtain the second modulation feature.

[0054] Steps 2-6: Input the second modulation feature into the feedforward fully connected layer for feature mapping, scale the mapped feature output by the feedforward fully connected layer, and obtain the output of the first layer diffusion transform block based on the scaled mapped feature and the scaled temporal correlation feature.

[0055] Steps 2-7: Use the output of the first-layer diffusion transform block as the input feature of the second-layer diffusion transform block, and obtain the global time-frequency features based on the output of the second-layer diffusion transform block.

[0056] Specifically, the temporal feature representation and the first modulation feature are input into the cross-attention layer, and the temporal correlation feature is output, including steps 3-1 to 3-2: Step 3-1: Use the first modulation feature as the query vector, and the temporal feature representation as the key vector and value vector, and calculate the cross-attention weights.

[0057] Step 3-2: Use cross-attention weights to perform weighted aggregation of temporal feature representations to obtain temporal correlation features.

[0058] For example, cross-attention weights The calculation formula is: , in, For query vector, , These are the key vector and the value vector, respectively. In this process, the output features are composed of attention weights and... The weighted summation is obtained by calculating the similarity between the corresponding query in the frequency domain and all keys in the time domain. The fused features contain both time-domain dynamics and frequency-domain structural information, and the dimensions are consistent with the input, enabling the model to more accurately focus on the key characteristics of the original signal, thereby improving the noise prediction accuracy and overall performance.

[0059] Specifically, the diffusion transform block in the noise prediction unit designed in this application incorporates an adaptive adjustment mechanism on top of the standard layer normalization, enabling the subnetwork to adaptively adjust according to input features. This breaks through the constraints of fixed normalization parameters, thus more flexibly adapting to diverse task requirements and data distribution characteristics. Simultaneously, conditional information such as time step embeddings and true fault category labels attached to the input sequence are fused using contextual conditionalization, and a cross-attention branch is added to the multi-head self-attention mechanism to enhance the utilization efficiency of conditional information and feature extraction performance. Furthermore, the parameters of specific linear mappings such as the conditional modulation layer and the output layer are initialized to zero to improve numerical stability in the early stages of training. Each layer can perform feature modulation through learnable scaling and translation operations, and a scaling regression operation is introduced at the end of the residual connection to further ensure the stability of the training process.

[0060] In the Diffusion Transformer architecture, the noise prediction unit uses attention mechanism and residual connection to realize the dynamic allocation of cross-layer feature weights, which enables the model to effectively capture the intrinsic relationship between input sequence and conditional information. This design not only improves the extraction accuracy of context features, but also further enhances the diversity and expressive power of latent feature representation, thereby effectively reducing the risk of overfitting the model on training data and improving the overall generalization performance.

[0061] First, the real fault category label serves as a semantic condition to guide the model in generating samples with specific fault characteristics. The time step represents the current stage of the diffusion process. After being parameterized into a learnable semantic vector, it participates in model training. Its value reflects the noise intensity level at that moment, thereby assisting the model in dynamically adjusting the noise prediction strategy according to the diffusion process. The two types of conditional information are encoded and weighted by the embedding mechanism, and then injected into the parameter feature modulation in the DiT Block.

[0062] Traditional layer normalization uses statically set scaling and translation parameters, which lack conditional awareness and are difficult to adjust for input features, category attributes, or diffusion stages. To address this, this application designs an adaptive normalization layer that uses a lightweight multilayer perceptron to dynamically obtain scaling and translation factors in a fine-grained manner. This enables the deep integration of conditional information into the generation process, avoids gradient anomalies caused by static normalization, and ensures that the model can perform stable and controllable high-quality noise prediction in the latent space.

[0063] Secondly, the self-attention mechanism can effectively capture global dependencies and improve feature representation capabilities. Building on this, this application further introduces a cross-attention mechanism to aggregate feature information from the time domain and frequency domain, constructing a richer fusion representation. This mechanism is derived from the self-attention idea, with the core difference being the different sources of the query, key, and value: in this application, the query is taken from frequency domain variables, while the key and value originate from time domain signals. This design enables explicit interaction between time-domain dynamic characteristics and frequency-domain structural information, enhancing the model's comprehensive ability to characterize the original signal.

[0064] Furthermore, such as Figure 4 The diagram shown illustrates the process of generating diffusion-generated samples according to this application. After obtaining the trained diffusion generation module, it further includes: The target quantization latent representation output by the trained diffusion generation module is input into the codebook quantization unit in the trained vector quantization module. The Euclidean distance between the target quantization latent representation and all codewords in the learnable codebook is calculated. The quantization latent variable of the target quantization latent representation is obtained based on the codeword index with the smallest Euclidean distance. The quantized latent variables are input into the decoder in the trained vector quantization module for decoding, and the output is diffused to generate samples.

[0065] This application also evaluates the samples generated by the above method and the corresponding fault diagnosis results: 1. MK-MMD Quantitative Assessment An effective generative model needs to capture the local and global features of real data sufficiently, so that the distribution of generated samples fits the initial data distribution as closely as possible. To this end, this application uses Multi-kernel Maximum Mean Discrepancy (MK-MMD) as an evaluation metric. It quantifies the degree of fit between the distributions of generated samples and real samples by calculating the mean difference in the mappings between their respective feature spaces. The MK-MMD values ​​of the generated samples for each model are shown in Table 1. Table 1

[0066] As shown in Table 1, the VQ-DiT method provided in this application exhibits significantly lower MK-MMD values ​​than baseline methods such as SMOTE, CVAE, CGAN, and DDPM under a full load condition of 0-3hp. The average MK-MMD value across all load conditions is only 0.0805, demonstrating extremely strong distribution alignment capabilities. Taking the most demanding 3hp load condition as an example, the MK-MMD value of VQ-DiT is as low as 0.089, which is 51.9% lower than SMOTE (0.185), 42.9% lower than CVAE (0.156), 71.5% lower than CGAN (0.312), and 79.2% lower than DDPM (0.428) compared to other methods. The results demonstrate that the samples generated by VQ-DiT have the smallest distribution difference in the feature space compared to the real samples, and the generation quality far exceeds all comparison baselines.

[0067] 2. Fault Diagnosis Results To comprehensively evaluate the mechanical fault diagnosis performance of various generative methods, comparative fault diagnosis experiments were conducted under full load conditions (T0-T3) on three datasets: SCARA, SDUST, and CWRU. The experiments compared diffusion-based methods such as VQDiT-CA, DiT (U-Net), and DDPM, as well as several GAN-based baseline methods. A blank control group (None) was also included. The diagnostic accuracy of each method is shown below. Figure 5 As shown.

[0068] The experimental results show that the VQ-DiT method exhibits the best fault diagnosis performance on all three datasets, with a cross-dataset average accuracy of 97.7%, which is significantly better than all other comparative methods.

[0069] On the CWRU dataset, VQ-DiT achieves an accuracy of 98.3%-99.0% across all scenarios, making it the best performing method on this dataset. It significantly outperforms similar diffusion methods such as DiT (U-Net) and DDPM, and also far surpasses various GAN-type methods.

[0070] On the SCARA dataset, VQ-DiT maintained an accuracy of over 96.5% across all scenarios, ranking first in overall performance. Other diffusion-type methods followed closely behind, while most GAN-type methods showed a significant performance decline.

[0071] On the SDUST dataset, VQ-DiT achieved an accuracy of over 96% across all scenarios, maintaining an absolute leading advantage. Other diffusion-type methods maintained stable performance, while GAN-type methods experienced severe performance degradation, with the lowest accuracy reaching only 77.2%.

[0072] It is worth noting that diffusion-based methods maintained stable diagnostic performance across all datasets and operating conditions, without significant performance fluctuations. In contrast, GAN-based methods exhibited vastly different performance across different datasets and operating conditions, demonstrating significantly inferior generalization ability and robustness compared to diffusion-based methods. The blank control group showed no effective diagnostic capability, further validating the effectiveness of the aforementioned deep learning methods. Overall, VQ-DiT achieved optimal diagnostic accuracy across all test scenarios, fully demonstrating its superior performance and strong robustness in mechanical fault diagnosis tasks.

[0073] In summary, the VQ-DiT model can generate synthetic samples that are highly consistent with the distribution of real samples. These generated samples not only closely resemble real samples in the feature space but also maintain good class separability, providing reliable data augmentation support for fault diagnosis tasks on imbalanced datasets.

[0074] This application designs a three-level data augmentation architecture for rotating machinery fault diagnosis, consisting of time-frequency representation, vector quantization, and diffusion generation. This architecture converts one-dimensional vibration time-series signals into two-dimensional time-frequency representations through short-time Fourier transform. After being mapped to a discrete latent space by a vector quantization module, the latent space diffusion generation and decoding reconstruction are completed based on a Diffusion Transformer architecture with time-frequency cross-attention. This enables the directional and high-quality generation of minority class samples of faults and completes the dataset equalization process.

[0075] Meanwhile, a two-stage model training and sample generation method is adopted. This method first pre-trains a vector quantization autoencoder using source domain fault data to complete the latent space mapping modeling of fault signals; then, it fixes the parameters of the vector quantization module and trains a DiT-based noise prediction network; finally, it generates discrete latent variables through iterative denoising, which are then restored as fault samples by the decoder. This solves the problem of mutual interference between latent space mapping and diffusion generation processes during end-to-end training and ensures the stability of model training.

[0076] Furthermore, a vector quantization autoencoder structure adapted to rotating machinery fault signals was designed. This structure consists of an encoder with residual skip connections, a learnable codebook, and a decoder. Deep time-frequency features are extracted through downsampling convolution and residual modules, compressed into a continuous latent space representation, and mapped from the continuous latent space to the discrete latent space using a nearest neighbor search strategy. This significantly reduces data dimensionality and computational overhead of subsequent diffusion models while fully preserving the key impact and modulation features of the fault signal. Additionally, an end-to-end training method for vector quantization based on a pass-through estimator was designed. This method constructs an approximate gradient propagation path using a stopping gradient operator, solving the technical problem of non-differentiability in nearest neighbor quantization operations and achieving collaborative optimization of the encoder, codebook, and decoder, as well as end-to-end model training. Furthermore, this application proposes a multi-loss joint optimization strategy for the vector quantization module. This strategy ensures sample reconstruction quality through reconstruction loss, achieves adaptive learning of codebook vectors through embedding loss, and constrains the fluctuation amplitude of encoder output through commitment loss. The weighted joint optimization of these three losses improves the stability of latent space modeling and the ability to represent fault features.

[0077] In addition, this application designs a DiT Block-based diffusion model noise prediction architecture. This architecture adopts a two-layer DiT Block cascaded structure and designs an adaptive normalization layer. It dynamically generates scaling and translation factors through a lightweight multilayer perceptron, deeply integrating time step embedding and fault category label conditional information into the noise prediction process. This achieves targeted generation of specific fault categories, breaking through the parameter constraints of traditional static normalization and the global feature extraction bottleneck of the U-Net architecture. Simultaneously, a time-frequency domain cross-attention mechanism for fault signals is designed. This mechanism uses the frequency domain latent variable token as the query and the time-domain features extracted from the original one-dimensional vibration signal mapping by a lightweight convolutional coding network as the key and value. This achieves explicit interaction and deep fusion of time-domain dynamic characteristics and frequency-domain structural information, enhancing the model's ability to mine time-frequency domain correlation features of fault signals and improving noise prediction accuracy and the realism of generated samples. In addition, a diffusion preprocessing and conditional embedding method adapted to discrete latent variables is proposed. This method divides the discrete latent variables output by vector quantization into patch sequences and completes linear mapping. It combines learnable position encoding, time step embedding and category conditional embedding to adapt to the sequence input requirements of Transformer. At the same time, a Cosine noise scheduling strategy is used to complete the noise processing of latent variables, ensuring the stability and controllability of the diffusion generation process.

[0078] In summary, by using the minority class fault samples generated by the method provided in this application to supplement the original imbalanced dataset, a balanced training dataset can be constructed, providing sufficient and high-quality data support for subsequent intelligent fault diagnosis models, and significantly improving the fault classification accuracy and generalization performance of the diagnostic models.

[0079] Those skilled in the art will understand that embodiments of this application can be provided as methods, systems, or computer program products. Therefore, this application can take the form of a completely hardware embodiment, a completely software embodiment, or an embodiment combining software and hardware aspects. Furthermore, this application can take the form of a computer program product embodied on one or more computer-usable storage media (including but not limited to disk storage, CD-ROM, optical storage, etc.) containing computer-usable program code.

[0080] This application is described with reference to flowchart illustrations and / or block diagrams of methods, apparatus (systems), and computer program products according to embodiments of this application. It will be understood that each block of the flowchart illustrations and / or block diagrams, and combinations of blocks in the flowchart illustrations and / or block diagrams, can be implemented by computer program instructions. These computer program instructions can be provided to a processor of a general-purpose computer, special-purpose computer, embedded processor, or other programmable data processing apparatus to produce a machine, such that the instructions, which execute via the processor of the computer or other programmable data processing apparatus, generate instructions for implementing the flowchart... Figure 1 One or more processes and / or boxes Figure 1 A device that provides the functions specified in one or more boxes.

[0081] These computer program instructions may also be stored in a computer-readable storage medium that can direct a computer or other programmable data processing device to function in a particular manner, such that the instructions stored in the computer-readable storage medium produce an article of manufacture including instruction means, which are implemented in a process Figure 1 One or more processes and / or boxes Figure 1 The function specified in one or more boxes.

[0082] These computer program instructions may also be loaded onto a computer or other programmable data processing equipment to cause a series of operational steps to be performed on the computer or other programmable equipment to produce a computer-implemented process, thereby providing instructions that execute on the computer or other programmable equipment for implementing the process. Figure 1 One or more processes and / or boxes Figure 1 The steps of the function specified in one or more boxes.

[0083] Obviously, the above embodiments are merely illustrative examples for clear explanation and are not intended to limit the implementation. Those skilled in the art will recognize that other variations or modifications can be made based on the above description. It is neither necessary nor possible to exhaustively list all possible implementations here. However, obvious variations or modifications derived therefrom are still within the scope of protection of this invention.

Claims

1. A vector quantization diffusion method for data augmentation in mechanical fault diagnosis, characterized in that, include: Obtain training samples from the mechanical fault diagnosis training set and convert the training samples into a frequency domain representation; The frequency domain representation is encoded by the encoder in the vector quantization module, and the output is a continuous latent variable. The continuous latent variable is input into the codebook quantization unit in the vector quantization module. The Euclidean distance between the continuous latent variable and all codewords in the learnable codebook is calculated. The quantized latent representation of the continuous latent variable is obtained based on the codeword index with the smallest Euclidean distance. The quantized latent representation is input into the decoder in the vector quantization module for decoding, and the reconstructed sample is output. The reconstruction loss is constructed based on the errors of training samples and reconstructed samples; the codebook embedding loss is constructed based on the errors of the stopping gradient terms of quantized latent representation and continuous latent variables; the commitment loss is constructed based on the errors of the stopping gradient terms of continuous latent variables and quantized latent representation; and the codebook utilization regularization term is constructed based on the usage frequency of each codeword in the learnable codebook. A quantization loss function is constructed based on reconstruction loss, codebook embedding loss, commitment loss, and codebook utilization regularization term, thereby iteratively training the vector quantization module and using the trained vector quantization module to reconstruct the training samples. The iterative training of the vector quantization module also includes: The training samples are input into the encoder in the trained vector quantization module for encoding, and the output is a continuous latent variable. The continuous latent variable is input into the codebook quantization unit in the trained vector quantization module. The Euclidean distance between the continuous latent variable and all codewords in the learnable codebook is calculated. The quantized latent representation of the continuous latent variable is obtained based on the codeword index with the smallest Euclidean distance. The quantized latent representation is input into the noise addition unit in the diffusion generation module to add noise, resulting in a noisy representation. The Transformer encoder in the noisy representation input diffusion generation module is modeled to convert the noisy representation into a noisy sequence. The noise-added sequence, the true fault category of the training samples, and the time step corresponding to the diffusion step size are input into the noise prediction unit in the diffusion generation module for noise prediction and denoising. The noise prediction error is used to construct a noise prediction loss function, thereby iteratively training the diffusion generation module to obtain a trained diffusion generation module for denoising the quantized latent representation and outputting the target quantized latent representation.

2. The vector quantization diffusion method for enhancing mechanical fault diagnosis data according to claim 1, characterized in that, When iteratively training the vector quantization module, the pass-through estimator is used to perform quantization mapping during forward propagation and to construct an encoder input approximate gradient channel during backpropagation, thereby using the pass-through approximation method to pass the gradient and perform backpropagation gradient update.

3. The vector quantization diffusion method for enhancing mechanical fault diagnosis data according to claim 1, characterized in that, Continuous latent variables Represented as: ， in, , , , and These represent the height and width of the frequency domain representation, respectively. and This represents the height and width of the frequency domain representation after downsampling. Indicates the number of downsampling attempts; The encoder consists of multiple convolutional layers and residual blocks, and compresses the spatial dimension through downsampling operations. The frequency domain representation of the training samples; Quantitative latent representation for: ， ， in, Represents the position in the quantized latent representation The quantization value at that location; This represents the k-th codeword in the learnable codebook; This indicates the number of codewords in the learnable codebook; Represents position in continuous latent variables The value of the variable at that location; This represents the codeword index with the smallest Euclidean distance.

4. The vector quantization diffusion method for enhancing mechanical fault diagnosis data according to claim 3, characterized in that, Quantization loss function Represented as: ， ， ， ， ， in, Indicates the reconstruction loss; Indicates training samples; Indicates the reconstructed sample; Indicates codebook embedding loss; Represents the stopping gradient term of a continuous latent variable. This indicates the stopping gradient operator during forward propagation. During back propagation ; express The weights; Indicates a commitment to loss; This represents the stopping gradient term of the quantized latent representation; express The weights; Regular expressions representing codebook utilization; This represents the frequency of use of the k-th codeword in the learnable codebook; express The weight.

5. The vector quantization diffusion method for enhancing mechanical fault diagnosis data according to claim 1, characterized in that, Noise addition representation for: ， in, This represents the cumulative decay coefficient from the initial time to step t; This represents real noise and conforms to a Gaussian distribution. This represents the quantized latent representation output by the codebook quantization unit in the trained vector quantization module; Noise prediction loss function Represented as: ， in, Represents the expectation of a fully random sample; Indicates predicted noise; Indicates a time step; Indicates the actual fault category.

6. The vector quantization diffusion method for enhancing mechanical fault diagnosis data according to claim 1, characterized in that, The noisy sequence, the true fault category of the training samples, and the time step corresponding to the diffusion step size are input into the noise prediction unit in the diffusion generation module for noise prediction and denoising, including: The noisy sequence is sequentially subjected to patch segmentation and linear transformation, feature recombination and linear transformation to obtain serialized input features; The true fault categories of the training samples are input into the label embedding layer and mapped to obtain the category conditional embedding vector. The diffusion time step is then processed sequentially through temporal embedding, linear transformation, ReLU activation, and quadratic linear mapping to generate temporal features. The category conditional embedding vector and the temporal features are fused to obtain the conditional embedding features. The serialized input features and conditional embedding features are input into two concatenated diffusion transform blocks to obtain global time-frequency features; The global time-frequency features are sequentially processed through layer normalization, linear dimension transformation, and feature recombination to obtain the noise corresponding to the current time step for single-step noise estimation. The noise is then gradually removed by iterative multi-time-step reverse diffusion.

7. The vector quantization diffusion method for enhancing mechanical fault diagnosis data according to claim 6, characterized in that, The serialized input features and conditional embedding features are input into a two-layer concatenated diffusion transform block to obtain global time-frequency features, including: The serialized input features are used as the input features of the first layer diffusion transform block, and the first adaptive layer normalization is used to normalize the input features to obtain the first normalized features. The first normalized feature is scaled and shifted and modulated using the linearly transformed conditional embedding feature to obtain the first modulated feature; The original time-domain signal of the training samples is mapped to a time-domain feature representation with the same dimension as the quantized latent representation. The time-domain feature representation and the first modulation feature are input into the cross-attention layer, and the output is the time-series correlation feature. The scaled temporal correlation features and input features are adaptively normalized using the second adaptive layer normalization to obtain the second normalized features; The second normalized feature is scaled and shifted to obtain the second modulation feature; The second modulation feature is input into the feedforward fully connected layer for feature mapping. The mapped feature output by the feedforward fully connected layer is scaled. The output of the first layer diffusion transform block is obtained based on the scaled mapped feature and the scaled temporal correlation feature. The output of the first-layer diffusion transform block is used as the input feature of the second-layer diffusion transform block, and the global time-frequency features are obtained based on the output of the second-layer diffusion transform block.

8. The vector quantization diffusion method for enhancing mechanical fault diagnosis data according to claim 7, characterized in that, The temporal feature representation and the first modulation feature are input into the cross-attention layer, and the output temporal correlation features include: The first modulation feature is used as the query vector, and the temporal feature representation is used as the key vector and value vector to calculate the cross-attention weights. By using cross-attention weights to weight and aggregate the temporal feature representations, temporal correlation features are obtained.

9. The vector quantization diffusion method for enhancing mechanical fault diagnosis data according to claim 1, characterized in that, After obtaining the trained diffusion generation module, it also includes: The target quantization latent representation output by the trained diffusion generation module is input into the codebook quantization unit in the trained vector quantization module. The Euclidean distance between the target quantization latent representation and all codewords in the learnable codebook is calculated. The quantization latent variable of the target quantization latent representation is obtained based on the codeword index with the smallest Euclidean distance. The latent variables are quantized and input into the decoder in the trained vector quantization module for decoding, and the output is diffused to generate samples.

Images