Efficient sparse diffusion models

Abstract

A method comprises obtaining a pre-trained diffusion model trained on a dataset; performing hierarchical pruning on the pre-trained diffusion model to obtain a pruned model; performing fine-tuning on the pruned model; compiling the pruned model using a TensorRT optimizer after the fine-tuning to obtain an optimized model; and deploying, using a sparse storage format, the optimized model on a sparse tensor core in a graphical processing unit (GPU) for executing an accelerated inference.

Description

Efficient Sparse Diffusion ModelsCROSS-REFERENCE TO RELATED APPLICATIONS

[0001] This application claims priority to U.S. provisional patent application number 63 / 720,575 filed November 14, 2024, by Futurewei Technologies, Inc., and titled “Efficient Sparse Diffusion Models,” which is incorporated by reference in its entirety.TECHNICAL FIELD

[0002] The present application generally relates to generative models, and more specifically to efficient sparse techniques for optimizing diffusion models.BACKGROUND

[0003] Diffusion models (DMs) have been extensively used in image or video generation tasks and are recognized as one of the most popular generative models. However, their long inference time, and requirements for large amounts of memory limit their practical usability in real-world applications. These challenges arise from the iterative nature of diffusion models, where a sequence of forward and reverse diffusion processes is computed, often requiring hundreds or thousands of steps to generate a high-quality output.SUMMARY

[0004] The present disclosure provides efficient sparse techniques in diffusion models to enhance computational efficiency without compromising model performance. The efficient sparse techniques refer to methods that reduce the number of active (non-zero or significant) parameters, operations, or memory accesses in the diffusion models by eliminating redundant or low-importance elements. Embodiments of this disclosure disclose optimizing diffusion models by reducing the number of inference steps and eliminating low-importance model weights with a new pruning algorithm, exploring tensor runtime (TensorRT) optimization technique on a diffusion model, designing efficient memory usage strategies with new sparse storage format, and leveraging hardware accelerators like graphical processing units (GPUs) with specialized cores such as Sparse Tensor Cores (SPTC). These optimizations make diffusion models more practical for deployment in environments with limited computational resources or stringent latency requirements.

[0005] A first aspect relates to a method, the method comprising obtaining a pre-trained diffusion model trained on a dataset; performing hierarchical pruning on the pre-trained diffusion model to obtain a pruned model; performing fine-tuning on the pruned model; compiling the pruned model using a TensorRT optimizer after the fine-tuning to obtain an optimized model; and deploying, using a sparse storage format, the optimized model on a sparse tensor core in a GPU for executing an accelerated inference.

[0006] Optionally, in a second implementation according to any of the first aspect or any implementation thereof, wherein performing the hierarchical pruning comprises performing step pruning on the pre-trained diffusion model by: identifying and pruning early timesteps in a reverse denoising process based on a first threshold; and pruning adjacent timesteps in each pair of connected timesteps in the reverse denoising process; acquiring a loss function of each of remaining timesteps to be pruned after the step pruning; performing a Taylor expansion on the loss function associated with each of the remaining timesteps to obtain an importance score for each weight at the remaining timesteps; pruning weights corresponding to the remaining steps when corresponding importance scores fall below a preset threshold to obtain pruned weights; and applying a 2:4 pruning scheme on the pruned weights to obtain the pruned model.

[0007] Optionally, in a third implementation according to any of the first aspect or any implementation thereof, wherein the first threshold is applied to reduce weights of the early timesteps corresponding to low-level details.

[0008] Optionally, in a fourth implementation according to any of the first aspect or any implementation thereof, wherein the adjacent timesteps are pruned based on temporal locality.

[0009] Optionally, in a fifth implementation according to any of the first aspect or any implementation thereof, the method further comprises applying the 2:4 pruning scheme by selecting pruned weights in groups of four consecutive weights, wherein two weights of every four consecutive weights in each group are set to zero based on weight magnitude.

[0010] Optionally, in a sixth implementation according to any of the first aspect or any implementation thereof, wherein before applying the 2:4 pruning scheme, the method further comprises includes reordering the weights in each group of four consecutive weights.

[0011] Optionally, in a seventh implementation according to any of the first aspect or any implementation thereof, the method further comprises using a permutation pruning technique to maximize the weight magnitude of reordered weights.

[0012] Optionally, in an eighth implementation according to any of the first aspect or any implementation thereof, the method further comprises calculating the importance score for each weight using a combination of weight values and gradients.

[0013] Optionally, in a ninth implementation according to any of the first aspect or any implementation thereof, wherein compiling the pruned model using the TensorRT optimizer comprises: applying tensor fusion to combine multiple tensor operations in the pruned model into a single operation; and applying FP8 precision calibration to allow the pruned model to run at 8-bit floating point precision.

[0014] Optionally, in a tenth implementation according to any of the first aspect or any implementation thereof, wherein the sparse storage format comprises a Harwell Boeing-based Elliptical Packed (HB-ELLpack or HB-ELL) format to store and process a sparse matrix, and wherein applying the sparse storage format comprises: partitioning the sparse matrix into a plurality of blocks; identifying non-zero blocks of the plurality of blocks having at least one non-zero element; and storing the non-zero blocks using the HB-ELL format.

[0015] Optionally, in an eleventh implementation according to any of the first aspect or any implementation thereof, the method further comprises determining a number of ELL columns for each block based on the number of non-zero elements; and storing the number of ELL columns in a block column size (bcols) array to reduce zero-filling.

[0016] Optionally, in a twelfth implementation according to any of the first aspect or any implementation thereof, wherein the sparse matrix is partitioned into 2x2 or 4x4 sub-blocks, and wherein each block is processed in row-major order for spatial coherence.

[0017] Optionally, in a thirteenth implementation according to any of the first aspect or any implementation thereof, wherein each of the non-zero blocks comprises a value array storing nonzero elements of each block; and an offset array storing column indices of the non-zero elements.

[0018] Optionally, in a fourteenth implementation according to any of the first aspect or any implementation thereof, the method further comprises using a browvalptr array to locate a starting position of data for each block within the value array.

[0019] Optionally, in a fifteenth implementation according to any of the first aspect or any implementation thereof, the method further comprises performing a sparse matrix-matrix multiplication (SpMM) using the HB-ELL format.

[0020] Optionally, in a sixteenth implementation according to any of the first aspect or any implementation thereof, wherein TensorRT optimizer contributes a 33.5% speedup, the hierarchical pruning algorithm contributes a 34.2% speedup, and the sparse storage format contributes a 32.3% speedup in the optimized model.

[0021] Optionally, in a seventieth implementation according to any of the first aspect or any implementation thereof, wherein the combination of hierarchical pruning, TensorRT optimization, and sparse tensor core acceleration provides 3.7* speedup with TensorRT and 7.4* speedup with the sparse tensor core on an ImageNet dataset.

[0022] Optionally, in an eighteenth implementation according to any of the first aspect or any implementation thereof, wherein the pruned model is fine-tuned for 4 epochs using a scaled learning rate.

[0023] A second aspect relates to a device comprising a memory or storage means configured to store instructions; and one or more processors or processing means coupled to the memory or the storage means and configured to execute the instructions to cause the device to perform the method according to any of the preceding aspect or any implementation thereof.

[0024] A third aspect relates to a computer program product comprising computer-executable instructions stored on a non-transitory computer-readable storage medium, the computer-executable instructions when executed by a processor of a device, cause the device to perform the method according to any of the preceding aspect or any implementation thereof.

[0025] For clarity, any one of the foregoing embodiments may be combined with any one or more of the other foregoing embodiments to create a new embodiment within the scope of the present disclosure.

[0026] These and other features, and the advantages thereof, will be more clearly understood from the following detailed description taken in conjunction with the accompanying drawings and claims.BRIEF DESCRIPTION OF THE DRAWINGS

[0027] For a more complete understanding of this disclosure, reference is now made to the following brief description, taken in connection with the accompanying drawings and detailed description, wherein like reference numerals represent like parts.

[0028] FIG. 1 is a schematic diagram of 2:4 sparse matrix format and mapping to SPTC according to an embodiment of the present disclosure.

[0029] FIG. 2 is a schematic diagram of a diffusion model framework according to an embodiment of the present disclosure.

[0030] FIG. 3 is a schematic diagram of a hierarchical-pruning algorithm according to an embodiment of the present disclosure.

[0031] FIG. 4 is a schematic diagram of a HB-ELLpack format for an 8 x 8 sparse matrix using 2 x 2 blocks according to an embodiment of the present disclosure.

[0032] FIG. 5A-5B are schematic diagrams of images sampled on ImageNet datasets and Canadian Institute for Advanced Research (CIFAR)-10 datasets respectively according to an embodiment of the present disclosure.

[0033] FIG. 6 is a flowchart of an embodiment of a method according to an embodiment of the present disclosure.

[0034] FIG. 7 is a schematic diagram illustrating a network element according to an embodiment of the present disclosure.DETAILED DESCRIPTION

[0035] It should be understood at the outset that, although illustrative implementations of one or more embodiments are provided below, the disclosed systems and / or methods may be implemented using any number of techniques, whether currently known or in existence. The disclosure should in no way be limited to the illustrative implementations, drawings, and techniques illustrated below, including the exemplary designs and implementations illustrated and described herein, but may be modified within the scope of the appended claims along with their full scope of equivalents.

[0036] The present disclosure provides methods and apparatuses for efficient sparse techniques in diffusion models to enhance computational efficiency without compromising model performance. Embodiments of this disclosure includes optimizing diffusion models by reducing the number of inference steps and eliminating low-importance model weights with new pruning algorithm, exploring TensorRT optimization technique on diffusion model, designing efficient memory usage strategies with new sparse storage format, and leveraging hardware accelerators like GPUs with specialized cores such as Sparse Tensor Cores. These optimizations make diffusionmodels more practical for deployment in environments with limited computational resources or stringent latency requirements.

[0037] DMs have emerged as powerful tools in generative tasks, including image synthesis, natural language processing, and molecular generation for their ability to produce high-quality samples. These models typically operate in two stages: a forward diffusion process that perturbs the data distribution to learn time-dependent score functions, and a reverse sampling process that generates data samples iteratively from a prior distribution. This two-phase process enables DMs to generate highly realistic outputs for a wide range of applications. However, the computational demands of DMs, particularly in terms of memory and processing power, present significant challenges for their deployment, particularly in mobile and embedded devices where resources are limited.

[0038] To address the computational load during DM inference, techniques such as training- free samplers and distillation approaches have been proposed to reduce the number of inference steps. However, these techniques often lead to a trade-off in sample quality. In particular, the quality of generated samples suffers when too few inference steps are used, as the model is unable to accurately reconstruct the complex high-dimensional data space, such as images or videos. This challenge is further exacerbated on mobile devices, where limited computational power and memory severely constrain the model’s ability to perform the necessary number of operations to produce high-quality outputs. The existing pruning methods overlook the hardware features of modern GPUs or fail to produce efficient pruning configurations, resulting in suboptimal performance.

[0039] Furthermore, recent advancements in GPU architecture, such as NVIDIA’ s Ampere and Hopper GPUs, include support for fine-grained structured sparsity and advanced tensor core features. These architectures provide 2:4 sparse mode that enables a 50% sparsity rate by retaining only two out of every four adjacent weights in a pre-trained model. This approach facilitates faster matrix multiplications by focusing only on non-zero values, theoretically offering up to a 2x speedup in computational efficiency. For diffusion models, incorporating this 2:4 sparsity mode at each inference step can dramatically reduce the computational load, making it a valuable optimization for tasks that require iterative refinement, such as generative tasks. However, these architectures fail to account for the unique structural characteristics of these models or do not take full advantage of the GPU’s capabilities, particularly when it comes to specialized features like floating-point 8-bit (FP8) precision.

[0040] The present disclosure discloses a hierarchical pruning algorithm for diffusion models, combining the structural properties of these models with the Hopper GPU architecture’s sparse tensor core capabilities to integrate sparsity into diffusion models by leveraging new GPU hardware features, thereby optimizing inference efficiency. The disclosed embodiments integrate sparsity into the diffusion model by employing 2:4 structured sparsity and additional scaling techniques, thereby enhancing efficiency for real-world applications.

[0041] In an embodiment, a comparative experiment is performed using several DMs on NVIDIA®’ s Hl 00 GPU with datasets such as ImageNet® and CIFAR10. The evaluation focuses on Frechet Inception Distance (FID) to assess sample quality, and Multiply-Accumulate Operations (MACs) to measure computational efficiency. The experimental results demonstrate that the method disclosed in this application achieves significant efficiency improvements while maintaining or enhancing sample quality, making it promising solution for deployment in real-world scenarios with limited computational resources.

[0042] In an embodiment, the present disclosure provides an approach for accelerating DMs using TensorRT. This includes a three-layer hierarchical pruning algorithm designed for diffusion models, enabling structured sparsity and improved efficiency. Additionally, a sparse tensor storage format is disclosed to effectively bridge the gap between sparse diffusion models and GPU tensor cores, thereby facilitating high-performance inference on modern hardware.

[0043] Related Work

[0044] Diffusion models have seen rapid advancements, with foundational works such asDenoising Diffusion Probabilistic Models (DDPM) and Denoising Diffusion Implicit Models (DDIM) laying the groundwork for a wide range of applications. These models rely on a two-phase process — a forward diffusion process and a reverse denoising process, where the reverse step is trained to progressively remove noise from inputs. The existing methodologies to improve the efficiency of diffusion models mainly fall into three categories: network architecture optimization, training process refinement, and sampling acceleration. Many diffusion models employ a U-Net architecture as the core denoiser, and their efficiency can be improved by incorporating hierarchical designs or performing training within a latent space. Recent research also suggests for integrating more efficient layers or architectures within the U-Net denoiser to enhance performance, thereby enabling the model to learn higher-quality image representations during training.

[0045] Additionally, significant efforts are directed towards boosting training efficiency, with some studies showing that diffusion training can be accelerated by adjusting the weight distribution across timesteps. Training can also be optimized by applying diffusion models at the patch level. Some approaches focus on sampling efficiency, which often is achievable without retraining the model. Methods such as step reduction, early stopping, and knowledge distillation have proven effective in reducing the number of sampling iterations required, thereby cutting down inference time while maintaining output quality.

[0046] Model Pruning: The concept of model sparsity has gained traction, with techniques such as pruning, quantization, and low-rank factorization applied across different neural networks. Sparse models are not only computationally efficient but also often exhibit better generalization properties. Sparse model is particularly achieved through the application of network pruning techniques. Pruning methodologies generally fall into two main categories: structural pruning and unstructured pruning. Structural pruning distinguishes itself by physically removing parameters and substructures from networks, while unstructured pruning effectively masks parameters by setting them to zero.

[0047] Most network pruning research has focused on discriminative tasks, such as classification tasks. Few studies have explored the potential of pruning techniques in generative tasks, such as generative adversarial networks (GANs) compression. Recently, applying structural pruning techniques to Diffusion Models has introduced unique challenges, prompting a reevaluation of traditional approaches. While some methods have attempted to leverage modern GPU features, they often fall short by not accounting for the distinct architectural and operational characteristics of diffusion models, thereby necessitating the development of more tailored pruning strategies.

[0048] The present disclosure provides a dedicated pruning method that effectively combines diffusion model features with GPU hardware optimizations to achieve efficient, high-quality model inference. The present disclosure explores the intersection of diffusion models, sparsity, and GPU hardware capabilities.

[0049] FIG. 1 is a schematic diagram 100 of 2:4 structured sparse matrix and its mapping to sparse tensor cores (SPTC) according to an embodiment of the present disclosure. NVIDIA Hopper GPUs have expanded their tensor core units (TCUs) to support row-wise 2:4 sparsity, incorporating hardware capabilities for sparse computation. These enhanced TCUs, known as SPTCs, allow for efficient processing of sparsely structured data. To leverage SPTCs, the first operand in tensor operations is formatted in NVIDIA’ s N:M sparsity format, where N representsthe maximum number of non-zero elements in a block of M values. As shown in FIG. 1, a sparse R matrix array 102 is an uncompressed sparse matrix. A compressor module 104, represented in the figure by the arrow and word Compress, is configured to receive and compress the sparse matrix array 102 into a compressed data file. The compressed data file involves two main components: (1) an R * K / 2 matrix 106 that stores the non-zero values from the matrix array 102, and (2) a metadata structure 108 that encodes the positions of non-zero elements within each group of four values. The metadata structure 108 provides positional information that enables hardware to efficiently extract and align corresponding elements from a dense matrix B 110 during matrix multiply-accumulate (MMA) operations. The dense matrix B 110 is then mapped to an SPTC 112, allowing for the execution of 2:4 sparse operations within the SPTC architecture.

[0050] FIG. 2 is a schematic diagram of an optimized DM framework 200 according to an embodiment of the present disclosure. The framework 200 outlines the process of optimizing a pretrained dense DM for efficient inference on sparse GPUs, specifically targeting 2:4 Sparse Tensor Cores. In an embodiment, the framework 200 comprises a dataset 202 and a pretrained dense DM block 204. The pretrained dense DM block 204 represents as an initial point of the optimization process, where a pretrained dense diffusion model (DM) is used as the baseline. The pretrained model is trained on the dataset 202 and exhibits a high computational cost due to the large number of multiply-accumulate operations (MACs) and its dense structure. To reduce this computational burden, the framework 200 incorporates a hierarchical pruning algorithm block 206, which applies a structured three-layer pruning process. The pruning process is performed across three different layers of the model: (i) step layer pruning that removes redundant operations within the step layers of the diffusion process, (ii) U-Net layer pruning that reduces unnecessary operations and parameters within the U-Net architecture for denoising tasks, and (iii) 4-element layer pruning that targets further pruning in the 4-element layer, which is part of the network’s structure, designed to maintain critical features while reducing computational overhead. The hierarchical pruning exhibits reduced MACs, making the pruned model more efficient for sparse GPU inference.

[0051] Upon completion of the pruning process, the framework 200 proceeds to a fine-tuning block 208. This stage retrains the pruned model to recover any lost accuracy, adjusting weights and parameters to adapt to the sparsity while retaining performance. The output of this stage is the trained sparse DM block 210, which represents a sparse version of the original pretrained diffusion model. This model has been optimized through both pruning and fine-tuning, achieving a reduction incomputational complexity while maintaining high accuracy. The trained sparse DM block 210 proceeds for further optimization in the next stages of the framework. The framework 200 further comprises a compiling block 212, where the fine-tuned sparse model is compiled using TensorRT, an NVIDIA library designed for optimizing models for inference on GPUs. This step performs optimizations such as tensor fusion, which merges multiple operations to reduce memory access and improve speed, and FP8 quantization, which reduces numerical precision to 8-bit floating point to enhance efficiency with minimal accuracy loss.

[0052] Following compilation, the framework 200 further comprises a sparse storage format block 214. The sparse storage format block 214 ensures that the model’s sparse data is stored in an optimized format, specifically designed to deploy on 2:4 Sparse Tensor Cores for maximum acceleration. Finally, the optimized model is passed to a deployment block 216, where it is deployed to GPU hardware featuring 2:4 sparse tensor core architecture. The sparse tensor core architecture is designed to perform computations more efficiently on sparse data by taking advantage of the 2:4 sparsity pattern. The inference block 218 represents real-time inference using the optimized model, achieving both accelerated performance and energy efficiency. In essence, the framework 200 seeks to optimize a diffusion model by pruning and quantizing it, followed by leveraging efficient storage and hardware-specific accelerations for faster, more efficient inference on sparse GPUs.

[0053] Pruning algorithm

[0054] The present disclosure provides a hierarchical pruning algorithm designed to enhance model sparsity while preserving performance. This approach integrates three distinct pruning techniques such as step pruning, Taylor pruning, and 2:4 structured pruning into a unified framework. By combining these strategies, the algorithm achieves substantial reductions in active parameters with minimal impact on model accuracy. A feature of this method is the enforcement of a 2:4 sparsity constraint, whereby only two of every four consecutive parameters are retained, and the remaining two are set to zero. This structured sparsity pattern is particularly well-suited for execution on hardware that supports 2:4 Sparse Tensor Cores.

[0055] In an embodiment, assume the parameter 0 is represented as a two-dimensional (2D) matrix, where each sub-structure $i = [6io, 0ii, . . ., tk] corresponds to a row vector with K scalar parameters. Structural pruning aims to produce a sparse parameter matrix 9' that retains performance characteristics close to the original model. This is achieved by solving the following optimization problem:

[0056] Here, L(0) denotes the model’s loss function, ||0'||o represents the Lo norm, i.e., the number of non-zero elements (or row vectors) in the pruned parameter matrix, and s denotes a sparsity level of the pruned model. However, given the iterative nature of diffusion models, the training objective L can be seen as a composition of T interdependent tasks: {Zi, Z2, . . ., ZT}. These tasks are sequentially dependent, meaning that pruning-induced changes in one timestep can propagate and influence all subsequent steps. This interdependence makes pruning diffusion models fundamentally different and more complex than pruning models trained for single-task discriminative objectives.

[0057] First Layer: step pruning on T.

[0058] In an embodiment, the first layer of the hierarchical pruning algorithm focuses on step pruning, which is performed based on two key observations related to the behavior of diffusion models during the denoising process.

[0059] Observation 1 : The early steps in the denoising process tend to focus on reconstructing local features, such as textures, edges, and color information. In contrast, the later steps are more critical for capturing global semantics, such as object and shapes.

[0060] Observation 2: A strong temporal locality exists between adjacent steps in the denoising process, meaning that adjacent steps share similar characteristics. This property allows for joint pruning of adjacent steps without a significant drop in model performance.

[0061] To leverage these observations, the pruning process is tailored as follows. A low pruning threshold can be set for the initial denoising steps, reducing their importance using a heuristic algorithm to determine the number of early steps to retain. Additionally, for adjacent steps, direct pruning can be applied by targeting every other step, thereby reducing redundancy between these closely related steps. After applying step pruning, the impact of each remaining step can be evaluated and its associated loss component Lt in the pruning process, ensuring that the denoising procedure remains effective despite the reduced number of steps.

[0062] Second Layer: Taylor Expansion at loss component Lt

[0063] In an embodiment, the second layer of pruning involves the use of a Taylor expansion applied to the loss function Lt to estimate the impact of pruning individual weights on modelperformance. The Taylor series expansion of the loss around the original parameter set 0 is given by:Zt(0') = Zt (0) + VZt (0) • O' - 0) + 0(110' - 0II2)

[0064] Subtracting Zt(0) from both sides, the loss difference due to pruning becomes: £t (0') -£t (e) = V£t (9) • (0' - 9) + 0(110' - 9II2).

[0065] This linear approximation allows for an efficient estimation of how much each weight contributes to the model’s loss, facilitating the identification and removal of the least impactful weights. The importance score for each weight is computed using its magnitude and the gradient of the loss with respect to that weight. This enables the identification of weights whose removal would have minimal effect on the model’s predictive performance. Weights with low importance scores indicate a small contribution to Zt are selectively pruned. This results in a reduced parameter set while preserving the overall model accuracy.

[0066] Third Layer: 2:4 Pruning

[0067] The third layer of the hierarchical pruning algorithm introduces a 2:4 structured sparsity constraint, designed for efficient execution on hardware such as NVIDIA’ s Sparse Tensor Cores. In this layer, pruning is performed at a fine-grained level by operating on groups of four consecutive weights.

[0068] In an embodiment, for each group of four weights, the two weights with the lowest magnitudes are pruned, while the two with the highest magnitudes are retained. This approach ensures that the parameters contributing the most to the model’s representational power are preserved. To optimize the sparsity pattern, a technique known as permutation pruning is employed. The permutation pruning searches for an optimal reordering of weights such that, when the 2:4 pruning constraint is applied, the total magnitude of the retained weights is maximized. This strategy helps maintain a higher level of model expressiveness and accuracy, even after enforcing structured sparsity. By applying 2:4 pruning in this manner, the model achieves a balance between sparse representation and computational performance, enabling accelerated inference on compatible GPU architectures while minimizing performance degradation.

[0069] FIG. 3 is a schematic diagram of a hierarchical-pruning algorithm according to an embodiment of the present disclosure. In particular, FIG. 3 illustrates a sequence of actions (e.g.,lines 1-12) executed by the hierarchical pruning system to iteratively prune a pretrained diffusion model.

[0070] FIG. 4 is a schematic diagram illustrating an example of HB-ELL format applied to an 8 x 8 sparse matrix, partitioned using 2 x 2 blocks, according to an embodiment of the present disclosure. This format is designed to efficiently support hierarchical sparsity introduced through the multi-level pruning strategies described earlier. In an embodiment, the HB-ELL format is employed to store and process sparse matrix-dense matrix multiplications (SpMM) on Sparse Tensor Cores (STCs). The HB-ELL format is an integrated representation capable of supporting both element -wise sparsity and vector- wise sparsity.

[0071] As shown in FIG. 4, after executing the first layer (step pruning) and the second layer (Taylor pruning) on a pretrained diffusion model, the resulting structure can be represented as an 8 x 8 sparse matrix 402. The sparse matrix 402 comprises sparsified weight data resulting from the pruning process. The sparse matrix 402 is partitioned into multiple sub-blocks (e.g., 2 x 2 or 4 x 4), and only non-zero blocks — those comprising at least one non-zero element — are retained for subsequent operations. In the shown example, six non-zero blocks 404 remain within the 8 x 8 matrix. These non-zero blocks 404 may subsequently be compressed and organized into block rows 406 (e.g., browl to brow4), with each row comprising block groups of size 2. A 2:4 structured sparsity pattern is then applied to each block row to derive a 2:4 sparse matrix 408. In the resulting 2:4 sparse matrix 408, each 4 x 4 sub-block contains eight non-zero elements, where each row and column within the block has exactly two non-zero values. The retained non-zero entries are stored as sparse blocks 410 using the ELLPack format, allowing for regular, efficient computation. In an embodiment, each sparse block 410 is represented in ELL format using two arrays: a value array, which holds the non-zero weight values, and an offset array, which stores the position of each nonzero value within the block.

[0072] The number of ELL columns, denoted as Cs, may vary across the ELL representations of the blocks. For instance, three of the blocks have an ELL representation with only one column (i.e., C = 1), whereas two of the blocks include two ELL columns (i.e., C = 2). To manage this variability, an additional array, referred to as bcols, is maintained. This additional array records the number of ELL columns for each block. All blocks are stored in a row-major order (e.g., browl blocks are stored before brow2, and so on) to preserve spatial coherence. Using ELL format for small-sized blocks has advantages in storage efficiency and in reducing the range of offset indices,thereby enabling representation using fewer bits. For example, in the configuration shown in FIG. 4, only l-bit is used to represent the offset indices due to the restricted index range [0,1],

[0073] However, due to pruning operations and irregular sparsity patterns, the number of ELL columns (Cs) per block may not be uniform. Other ELL representations use padding to the maximum number of non-zeros per row, leading to zero-filling, which increases both memory usage and computational overhead. The disclosed HB-ELL format mitigates this inefficiency by supporting non-uniform column counts across blocks. Each matrix block retains metadata indicating its specific column count via the bcols array, thereby reducing the need for excessive zero-filling. This optimization results in decreased memory footprint and improved computation efficiency during sparse matrix operations.

[0074] Sparse Matrix-Matrix Multiplication (SpMM)s

[0075] SpMM is a used for computational kernel in neural network workloads. In an embodiment of the present disclosure, an SpMM algorithm is designed based on the HB-ELL format. The HB-ELL-based SpMM algorithm aligns with hardware tensor core operations, wherein each GPU warp is assigned to compute one dense output block of the resulting matrix C. In contrast to other dense matrix operations, the disclosed algorithm allows a warp to multiply a sequence of sparse blocks (from matrix A, stored in HB-ELL format) with corresponding dense blocks from matrix B. The results are then accumulated into matrix C.

[0076] In an embodiment, the HB-ELL-based SpMM algorithm leverages five auxiliary arrays in addition to the sparse value arrays (i.e., the value and offset arrays from ELL format) to represent the structure and layout of the sparse matrix A. These auxiliary arrays include as follows.

[0077] bcols: An array of size nb (number of blocks) that stores the number of ELL columns per sparse block and that indicates how many non-zero entries are stored per row within each sparse block.

[0078] bridx: An array of size nb that stores the starting row index in matrix A for each sparse block.

[0079] bcidx: An array of size nb that indicates the starting column index in matrix A for each block.

[0080] browvalptr: An array of size nb + 1 that points to the starting position of each block’s data in the value array of matrix A.

[0081] browptr: An array of size nbr + 1 (where nbr is the number of blocked rows) that points to the indices in the block list where each blocked row begins.

[0082] For illustrative purposes, and with reference to FIG. 4, consider a scenario where the number of non-zero blocks nb=6 and the number of blocked rows nbr=4. In this case, the browptr array is used to identify the range of block indices associated with each blocked row in matrix A. This array facilitates the traversal of sparse blocks during the matrix multiplication process by defining the start and end positions of each row segment in the compressed block list. During computation, the bridx and bcidx arrays are used to determine the actual output block row (cRow) and block column (bRow) positions in the output matrix C, considering that zero blocks or rows are omitted in the HB-ELL representation.

[0083] The browvalptr array provides direct access into the value array, allowing for efficient retrieval of the sparse values associated with a given block, without scanning through unnecessary entries. For each blocked row in the sparse matrix A, the corresponding dense blocks of matrix B are retrieved and multiplied, with the results accumulated into a warp-local register accumulator (ACCreg). The sparse values of matrix A are accessed using the offset provided by browvalptr. After completion of the multiplication and accumulation phase, the corresponding block of matrix C is loaded from memory. A scaling operation is applied using scalar coefficients a and 0, consistent with the standard matrix multiplication formula (e.g., C = aAB + 0C). The final results are written back to memory from a register-level output buffer (Creg), completing the update of the matrix C.

[0084] Implementation -Sparse Diffusion Pipeline

[0085] The present disclosure describes a pipeline for enabling structured sparsity in diffusion models at both training and inference stages. In an embodiment, the pipeline comprises three primary phases: sparsity induction phase, sparse training phase, and inference phase.

[0086] In the sparsity induction phase, structured pruning is applied to a pretrained diffusion model. Unlike unstructured pruning that zeros out individual weights, structured pruning removes entire computation units such as filters or neurons, based on their respective contributions to the overall model performance. The pruning criteria may be determined through various importance metrics, including, but not limited to, weight magnitude, Taylor expansion sensitivity, or gradientbased heuristics. This phase reduces model complexity and size while preserving its functionalintegrity. Following pruning, the model undergoes retraining with a focus on preserving the imposed sparsity pattern.

[0087] In the sparse training phase, the model is trained using a sparsity-aware optimization algorithm. The sparse training phase is configured to maintain the imposed structured sparsity by updating only the remaining (unpruned) parameters. In some embodiments, sparse training may employ gradient masking or reparameterization techniques to prevent updates to pruned weights, thereby reinforcing the sparsity structure during the optimization process. Once the model has been trained under structured sparsity, it enters the inference stage.

[0088] In the final inference phase, the trained sparse model is deployed. Due to the reduced number of active parameters and operations, the sparse model exhibits improved computational efficiency, resulting in faster inference times and lower memory usage. In an embodiment, the deployment phase leverages hardware-optimized sparse matrix-matrix multiplication (SpMM) routines, such as those enabled by the HB-ELL format, along with TensorRT optimizations including tensor fusion and FP8 precision calibration.

[0089] TensorRT Optimization

[0090] To maximize the runtime efficiency of the sparse diffusion model on modern GPU hardware, the present disclosure incorporates TensorRT optimization techniques, including tensor fusion and FP8 precision calibration. These are executed using the TensorRT inference engine or an associated TensorRT optimizer.

[0091] In an embodiment, tensor fusion is used to combine multiple tensor operations into a single GPU kernel launch. This optimization reduces the computational overhead associated with repeated memory accesses by minimizing the number of intermediate memory reads / writes. By allowing intermediate data to remain resident in the GPU’s high-speed memory (e.g., shared memory or registers), the data transfer latency between GPU cores is significantly reduced. This optimization is particularly beneficial for models with a large number of small-scale operations or layers, which is characteristic of the sparse diffusion model disclosed herein.

[0092] In an embodiment, the optimization pipeline leverages FP8 precision calibration, a form of post-training quantization that converts model weights and activations to the 8-bit floating point (FP8) format. FP8 is supported by advanced GPU architectures such as NVIDIA Hopper. Through FP8 calibration, models originally trained using higher-precision formats (e.g., FP16 or FP32) can be executed using FP8 arithmetic. This reduces both memory footprint and computationalcomplexity during inference, while maintaining an acceptable level of numerical accuracy. In an embodiment, FP8 calibration is applied post-training via quantization-aware calibration, enabling efficient deployment of the sparse diffusion model in resource-constrained or latency- sensitive environments. By integrating these TensorRT techniques — tensor fusion and FP8 quantization — the disclosed sparse diffusion model achieves optimized performance on GPU hardware, thereby supporting high-throughput inference while maintaining model fidelity.

[0093] Experiments

[0094] The present disclosure presents a comprehensive experimental evaluation of the disclosed sparse diffusion framework. The experiments are designed to showcase the performance improvements, computational efficiency, and sample quality retention achieved through the integration of structured sparsity, hierarchical pruning, and TensorRT optimizations. The evaluation compares the disclosed sparse diffusion model against other dense diffusion models across multiple datasets and benchmarking metrics, thereby demonstrating its effectiveness in real-world deployment scenarios.

[0095] Experimental Setup

[0096] In an embodiment, to assess the generalizability and effectiveness of the disclosed framework across diverse diffusion model architectures, experiments are conducted using the widely adopted pre-trained diffusion model, namely the Latent Diffusion Model (LDM). Furthermore, to validate compatibility with accelerated sampling strategies, the disclosed method is integrated with the 50-step Pseudo Linear Multistep Sampling (PLMS) technique, instead of employing the full 1000-step denoising process used in other diffusion-based generation.

[0097] In an embodiment, the experimental evaluation leverages two widely recognized image generation datasets: CIFAR-10 and ImageNet. For each dataset, the model generates 50,000 images, which are then evaluated to measure the quality of generated outputs. The Frechet Inception Distance (FID) score is used as the primary quantitative metric for assessing the performance of generative models. FID evaluates how closely the distribution of generated images matches that of real images, capturing both perceptual quality and diversity. Lower FID scores indicate better generation quality.

[0098] To benchmark the performance of the proposed sparse diffusion framework, it is compared against multiple state-of-the-art methods, including: (1) DeepCache, which is a diffusion acceleration technique that integrates pruning and caching mechanisms to improve generation speedwhile preserving quality, (2) Taylor Pruning, which is a gradient-based pruning technique that estimates the importance of weights using Taylor series approximations, and (3) automatic structured pruning (ASP), which is an optimization-driven approach for learning structured sparsity patterns during training. These comparisons allow for a rigorous assessment of the proposed framework’s efficiency, accuracy retention, and superiority in real-world deployment scenarios.

[0099] FIGS. 5 A and 5B are schematic diagrams of images sampled on ImageNet datasets and CIFAR-10 datasets according to an embodiment of the present disclosure. The experimental results demonstrate the effectiveness of the disclosed pruning and optimization techniques in preserving image quality while significantly reducing computational complexity. FIG. 5A illustrates the results 500A of applying hierarchical pruning to an LDM pre-trained on the ImageNet dataset at a 256 x 256 resolution. FIG. 5B illustrates similar results 500B for an LDM pre-trained on the CIFAR-10 dataset at a 32 x 32 resolution. The LDM architecture comprises an encoder, decoder, and a U-Net model-with the U-Net accounting for approximately 400 million parameters, while the autoencoder (encoder + decoder) has only 55 million parameters, the pruning effort primarily targets the U-Net. Consequently, the pruning process primarily focuses on the U-Net, as it represents the most computationally intensive portion of the model. The visual outcomes in FIGS. 5 A and 5B underscore the framework’s ability to maintain high fidelity of generated images, even after aggressive pruning, thereby delivering significant computational savings.

[0100] In an embodiment, a pruning threshold T = 0.1 is applied to exclude layers deemed to have sufficiently converged, thereby reducing computational cost during pruning. Using this threshold, the pruning process is limited to 534 denoising steps. Following this step-level reduction, importance metrics are calculated across all layers of the U-Net to guide the structured pruning process. A uniform channel sparsity rate of 30% is then enforced, ensuring consistent parameter reduction across layers. As a result of this hierarchical pruning strategy, the parameter count of the U-Net is reduced from approximately 400 million to 189.43 million, representing a substantial compression of the model. To recover any potential loss in performance due to pruning, the pruned model undergoes fine-tuning for 4 epochs using the official LDM training scripts. During fine-tuning, the learning rate is scaled down to 0.1 x the original base learning rate to enhance training stability and prevent overfitting or destabilization of the remaining weights.

[0101] Table 1 presents the quantitative results on the CIFAR-10 dataset. The results highlight the framework’s ability to maintain or improve generation quality, as measured by the FID, whileoperating under similar MACs constraints. Notably, the disclosed framework yields a more substantial acceleration ratio compared to DeepCache and other pruning strategies. Furthermore, heuristic evaluations are conducted to identify optimal pruning steps within the diffusion process. The pruning step numbers tested include N = 21, 22, 23, 25. The optimal setting is empirically selected based on the best trade-off between model performance (FID) and inference efficiency (MACs and latency).Table 1. CIFAR-10 32 x 32Method MACs J, Throughput f Speed T FID J,Taylor-Pruning 6.1G 9.81 l x 4.2ASP 6.1G 9.81 l x 4.17DeepCache 3.5G 13.46 1.38x 5.31Ours - N=21 4.16G 13.75 1.41 x 4.35Ours - N=22 3.55G 15.76 1.62x 4.70Ours - N=23 3.03G 18.13 1.86x 5.73Ours - N=25 2.64G 20.28 2.08x 9.75

[0102] Pruning Ratios

[0103] In an embodiment, the relationship between pruning ratios and model performance is systematically evaluated. Table 2 summarizes the results by reporting the number of MACs and the corresponding FID scores across varying levels of sparsity. As shown in Table 2, the experimental analysis reveals that diffusion models are sensitive to reductions in parameter count introduced through pruning. As the pruning ratio increases, a gradual degradation in generative quality is observed, highlighting a fundamental trade-off between computational efficiency and output fidelity. Specifically, when the pruning ratio is increased from 5% to 30%, the FID score deteriorates from 3.26 to 4.05. These findings suggest the importance of selecting an optimal pruning threshold that balances performance degradation with computational gain. In some embodiments, adaptive pruning strategies or layer-wise sparsity tuning may be applied to further mitigate quality loss while maintaining high acceleration.Table 2. Comparison of Pruning Techniques at Different Pruning RatesPruning Rate Method MACs J, FID J,Ours 147.6G 3.265% Taylor Pruning 1416.3G 3.32DeepCache 221.18G 5.32ASP 196.1G 5.88Ours 92.71G 4.0530% Taylor Pruning 588.1G 5.87DeepCache 147.4G 8.22ASP 119.9G 7.1Ours 71.96G 4.4250% Taylor Pruning 99.8G 9.16DeepCache 99.82G 12.26ASP 87.3G 7.91

[0104] TensorRT Optimization

[0105] In an embodiment, the performance of the disclosed sparse diffusion model is evaluated under different inference frameworks to quantify the benefit of both pruning and hardware-level optimizations. Table 3 presents the inference performance of a 50% pruned model on the ImageNet dataset (256 x 256 resolution). As shown in the evaluation, the pruning algorithm reduces the total number of MACs to 71.96 giga (G) MACs, resulting in an approximate 2* improvement in computational efficiency. When executed using the PyTorch framework, the pruned model achieves a 1.8 / speedup relative to the original dense model. Further acceleration is achieved through deployment using TensorRT, which yields a 3.7x speedup. These results highlight the synergistic impact of model pruning and TensorRT runtime optimization, enabling significant reductions in both computational cost and inference latency. This demonstrates the practical feasibility of deploying high-performance sparse diffusion models in real-time or resource-constrained environments.Table 3. Comparison of Pruning Techniques at Different Pruning Rates

[0106] Sparse Tensor Core Optimization

[0107] In an embodiment, the performance of the pruned diffusion model is further evaluated on specialized hardware acceleration platforms, including Sparse Tensor Cores (STC). Table 4 presents the inference speedup achieved by the 50% pruned model on both the ImageNet (256 x 256) and CIFAR-10 (32 x 32) datasets, using different frameworks and hardware configurations. Table 4 shows on the ImageNet dataset, the pruned model achieves a 3.7x speedup when executed using TensorRT, and a 7.4x speedup when accelerated using Sparse Tensor Cores. On the CIFAR- 10 dataset, TensorRT provides a 5.9x speedup, while Sparse Tensor Cores deliver an even greater improvement, achieving an I 1.5x speedup over the dense baseline. These results clearly demonstrate the advantages of combining structured model pruning with hardware-specific optimization techniques, such as Sparse Tensor Core utilization. The integration of these methods enables substantial reductions in inference time while maintaining effective generative performance, thus making the disclosed framework suitable for deployment on high-efficiency inference platforms.Table 4. Sparse Tensor Core Optimization

[0108] Performance Speedup Breakdown for Optimization Components

[0109] In an embodiment, the performance benefits of individual optimization components within the disclosed framework are quantified and analyzed. Table 5 presents a breakdown of the inference speedup contributed by three primary components: TensorRT inference engine, the structured pruning algorithm, and the sparse storage format (e.g., HB-ELL). The TensorRT component, which serves as an optimized runtime environment for deep learning inference, contributes a 33.5% performance speedup. This gain is attributed to improved kernel fusion, memory access patterns, and precision calibration optimizations provided by the TensorRT engine. The structured pruning algorithm achieves a 34.2% speedup by reducing the overall computational workload of the diffusion model. Through channel-level and filter-level pruning, unnecessary parameters and operations are removed, enabling more efficient model execution without significantly impacting generative performance. The sparse storage format, implemented using the HB-ELL scheme, contributes an additional 32.3% speedup. This gain results from reduced memory bandwidth usage and lower indexing overhead, allowing sparse data to be processed more efficiently on both general-purpose and specialized hardware (e.g., Sparse Tensor Cores). These results demonstrate that each optimization component independently contributes substantial performance improvements. When combined, the cumulative effect of these optimizations yields a highly efficient inference pipeline for sparse diffusion models, thereby validating the modular design of the disclosed framework.Table 5. Performance speedup for different components

[0110] Ablation Study[0U1] In an embodiment, an ablation study is conducted to investigate the impact of the early pruning step on the performance of the diffusion model. A heuristic approach is employed to determine the optimal step N at which pruning should be applied, as outlined in Algorithm 1. The evaluation spans a range of step values from N = 5 to N = 45, and the corresponding results aresummarized in Table 6. For each value of N, the number of MACs and the resulting FID score are recorded. A lower MAC count indicates greater computational efficiency, whereas a lower FID score corresponds to higher image generation quality. The results demonstrate a trade-off between computational efficiency and generative performance. As the number of steps N decreases, the MACs steadily decrease, reflecting reduced computational complexity. Initially, the FID scores remain stable; however, beyond a certain threshold, FID scores begin to increase, indicating a decline in image quality due to overly aggressive early pruning. Notably, the configuration corresponding to N = 20 steps is highlighted in Table 6 as offering an optimal balance. At this point, the model achieves significantly reduced computational load without a substantial degradation in image quality. This result suggests that selecting an early pruning step at or near N = 20 enables an effective compromise between runtime performance and output fidelity. This ablation study confirms the importance of step-wise pruning design and highlights the effectiveness of the proposed heuristic method for identifying optimal pruning intervals within diffusion-based generative models.Table 6. Performance metrics with different steps.

[0112] FIG. 6 illustrates a flowchart of an embodiment of a method 600 for optimizing a diffusion model, according to an embodiment of the present disclosure. In an embodiment, method600 is performed by at least one circuit, at least one system, at least one processor, at least one graphics processing unit, at least one parallel processor, and / or at least some other processor or component thereof described and / or shown herein.

[0113] In an embodiment, at a block 602, method 600 includes obtaining a pre-trained diffusion model trained on a dataset.

[0114] In an embodiment, at a block 604, method 600 includes performing hierarchical pruning on the pre-trained diffusion model to obtain a pruned model. As described above, the pretrained model undergoes hierarchical pruning to produce a sparsified version. The pruning algorithm integrates three key layers such as step pruning, Taylor pruning, and 2:4 pruning to achieve efficient sparsity in the model’s parameters while minimizing loss in model performance. The 2:4 pruning enforces a constraint that two out of every four consecutive parameters are set to zero, thereby reducing the number of active weights while maintaining the model’s output accuracy.

[0115] In an embodiment, at a block 606, method 600 includes performing fine-tuning on the pruned model. Once pruning is complete, the model is further fine-tuned. This step retrains the remaining weights to recover any potential accuracy loss and to adapt the model to the new sparse structure.

[0116] In an embodiment, at a block 608, method 600 includes compiling the pruned model using a TensorRT optimizer after the fine-tuning to obtain an optimized model. The fine-tuned model is compiled using NVIDIA TensorRT, a high-performance deep learning inference engine. TensorRT optimizes the model for GPU execution via tensor fusion and FP8 quantization

[0117] In an embodiment, at a block 610, method 600 includes deploying, using a sparse storage format, the optimized model on a sparse tensor core in a graphical processing unit (GPU) for executing an accelerated inference. The optimized model is deployed using a sparse storage format, such as HB-ELL, which is designed to align with the 2:4 sparsity structure. By leveraging the 2:4 sparsity pattern, the optimized model achieves higher inference speed and reduced energy consumption. Overall, the method 600 optimizes a diffusion model through hierarchical pruning and quantization, followed by efficient storage and hardware-specific acceleration for faster, more efficient inference on sparse GPUs.

[0118] FIG. 7 is a schematic diagram illustrating a network apparatus 700 according to an embodiment of the present disclosure. The network apparatus 700 is suitable for implementing thedisclosed embodiments as described herein. The network apparatus 700 comprises ingress ports / ingress means 710 (a.k.a., upstream ports) and receiver units (Rx) / receiving means 720 for receiving data; a processor, logic unit, or central processing unit (CPU) / processing means 730 to process the data; transmitter units (Tx)Ztransmitting means 740 and egress ports / egress means 750 (a.k.a., downstream ports) for transmitting the data; and a memory / memory means 760 for storing the data. The network apparatus 700 may also comprise optical-to-electrical (OE) components and electrical -to-optical (EO) components coupled to the ingress ports / ingress means 710, the receiver units / receiving means 720, the transmitter units / transmitting means 740, and the egress ports / egress means 750 for egress or ingress of optical or electrical signals.

[0119] The processor / processing means 730 is implemented by hardware and software. The processor / processing means 730 may be implemented as one or more CPU chips, cores (e.g., as a multi-core processor), field-programmable gate arrays (FPGAs), application specific integrated circuits (ASICs), and digital signal processors (DSPs). The processor / processing means 730 is in communication with the ingress ports / ingress means 710, receiver units / receiving means 720, transmitter units / transmitting means 740, egress ports / egress means 750, and memory / memory means 760. The processor / processing means 730 comprises a diffusion model module 770. The diffusion model module 770 is able to implement the methods disclosed herein. The inclusion of the diffusion model module 770 therefore provides a substantial improvement to the functionality of the network apparatus 700 and effects a transformation of the network apparatus 700 to a different state. Alternatively, the diffusion model module 770 is implemented as instructions stored in the memory / memory means 760 and executed by the processor / processing means 730.

[0120] The network apparatus 700 may also include input and / or output (I / O) devices or VO means 780 for communicating data to and from a user. The I / O devices or I / O means 780 may include output devices such as a display for displaying video data, speakers for outputting audio data, etc. The I / O devices or I / O means 780 may also include input devices, such as a keyboard, mouse, trackball, etc., and / or corresponding interfaces for interacting with such output devices.

[0121] The memory / memory means 760 comprises one or more disks, tape drives, and solid- state drives and may be used as an over-flow data storage device, to store programs when such programs are selected for execution, and to store instructions and data that are read during program execution. The memory / memory means 760 may be volatile and / or non-volatile and may be read-only memory (ROM), random access memory (RAM), ternary content-addressable memory (TCAM), and / or static random-access memory (SRAM).

[0122] While several embodiments have been provided in the present disclosure, it may be understood that the disclosed systems and methods might be embodied in many other specific forms without departing from the spirit or scope of the present disclosure. The present examples are to be considered as illustrative and not restrictive, and the disclosure is not to be limited to the details given herein. For example, the various elements or components may be combined or integrated in another system or certain features may be omitted, or not implemented.

[0123] In addition, techniques, systems, subsystems, and methods described and illustrated in the various embodiments as discrete or separate may be combined or integrated with other systems, modules, techniques, or methods without departing from the scope of the present disclosure. Other items shown or discussed as coupled or directly coupled or communicating with each other may be indirectly coupled or communicating through some interface, device, or intermediate component whether electrically, mechanically, or otherwise. Other examples of changes, substitutions, and alterations are ascertainable by one skilled in the art and may be made without departing from the spirit and scope disclosed herein.

Claims

CLAIMS1. A method comprising: obtaining a pre-trained diffusion model trained on a dataset; performing hierarchical pruning on the pre-trained diffusion model to obtain a pruned model; performing fine-tuning on the pruned model; compiling the pruned model using a TensorRT optimizer after the fine-tuning to obtain an optimized model; and deploying, using a sparse storage format, the optimized model on a sparse tensor core in a graphical processing unit (GPU) for executing an accelerated inference.

2. The method of claim 1, wherein the hierarchical pruning comprises a three-stage pruning process, wherein the three-stage pruning process comprises step pruning, a Taylor expansion-based pruning, and a 2:4 structured pruning, and wherein performing the hierarchical pruning comprises: performing the step pruning on the pre-trained diffusion model by: identifying and pruning early timesteps in a reverse denoising process based on a first threshold; and pruning adjacent timesteps in each pair of connected timesteps in the reverse denoising process; acquiring a loss function of each of remaining timesteps to be pruned after the step pruning; performing the Taylor expansion-based pruning on the loss function associated with each of the remaining timesteps to obtain an importance score for each weight at the remaining timesteps; pruning weights corresponding to the remaining steps when corresponding importance scores fall below a preset threshold to obtain pruned weights; and applying the 2:4 structured pruning on the pruned weights to obtain the pruned model.

3. The method of claim 2, wherein the first threshold is applied to reduce weights of the early timesteps corresponding to low-level details.

4. The method of claim 2, wherein the adjacent timesteps are pruned based on temporal locality.

5. The method of claim 2, further comprising applying the 2:4 structured pruning by selecting the pruned weights in groups of four consecutive weights, wherein two weights of every four consecutive weights in each group are set to zero based on weight magnitude.

6. The method of claim 5, wherein before applying the 2:4 structured pruning, the method further comprises reordering the weights in each group of four consecutive weights.

7. The method of claim 6, further comprising using a permutation pruning technique to maximize the weight magnitude of reordered weights.

8. The method of claim 2, further comprising calculating the importance score for each weight using a combination of weight values and gradients.

9. The method of claim 1, wherein compiling the pruned model using the TensorRT optimizer comprises: applying tensor fusion to combine multiple tensor operations in the pruned model into a single operation; and applying floating-point 8-bit (FP8) precision calibration to allow the pruned model to run at 8-bit floating point precision.

10. The method of claim 1, wherein the sparse storage format comprises a Harwell Boeingbased elliptical packed (HB-ELL) format to store and process a sparse matrix, and wherein using the sparse storage format comprises: partitioning the sparse matrix into a plurality of blocks; identifying non-zero blocks of the plurality of blocks having at least one non-zero element; and storing the non-zero blocks using the HB-ELL format.

11. The method of claim 10, further comprising: determining a number of ELL columns for each block based on the number of non-zero blocks; andstoring the number of ELL columns in a block column size (bcols) array to reduce zerofilling.

12. The method of claim 10, wherein the sparse matrix is partitioned into 2*2 or 4x4 sub-blocks, and wherein each block is processed in row-major order for spatial coherence.

13. The method of claim 10, wherein each of the non-zero blocks comprises: a value array storing non-zero elements of each block; and an offset array storing column indices of the non-zero elements.

14. The method of claim 13, further comprising using a browvalptr array to locate a starting position of data for each block within the value array.

15. The method of claim 10, further comprising performing a sparse matrix-matrix multiplication (SpMM) using the HB-ELL format on the sparse tensor core.

16. The method of claim 1, wherein TensorRT optimizer provides a 33.5% speedup, the hierarchical pruning algorithm provides a 34.2% speedup, and the sparse storage format provides a 32.3% speedup of the optimized model.

17. The method of claim 1, wherein a combination of the hierarchical pruning, TensorRT optimization, and sparse tensor core acceleration provides 3.7x speedup with TensorRT and 7.4x speedup with the sparse tensor core on an ImageNet dataset.

18. The method of claim 1, wherein the pruned model is fine-tuned for 4 epochs using a scaled learning rate.

19. A device comprising: a memory configured to store instructions; and one or more processors coupled to the memory and configured to execute the instructions to cause the device to perform a method according to any of claims 1-18.

20. A computer program product comprising computer-executable instructions stored on a non-transitory computer-readable storage medium, the computer-executable instructions when executed by one or more processor of a device, cause the device to perform a method according to any of claims 1-18.

Images