Method and apparatus for training a diffusion model to perform super-resolution

Abstract

The application describes a method for training a diffusion model to perform super-resolution of low or lower resolution images, which can be deployed on resource-constrained devices such as smartphon

Description

Field

[001] The present techniques generally relate to a method for training a diffusion model. In particular, the present techniques provide a method for training a diffusion model to perform super-resolution of low or lower resolution images, which can be deployed on resource-constrained devices such as smartphones. Background

[002] Latent diffusion models (LDMs) have shown remarkable progress in various image generation tasks due to their ability to leverage large-scale training to produce models with excellent generality. They have similarly been successfully applied to super-resolution, where the excellent image priors provided by SD can be readily exploited to obtain very naturallooking outputs. In other words, stable (latent) diffusion (SD) models have revolutionized several image generation tasks like text-to-image generation and editing, in-painting, video generation and super-resolution to name a few.

[003] However, their huge model size and computational complexity render their deployment on resource-limited edge devices like mobile phones prohibitive. Their high computation cost limits the more general deployability as the pre-existing computational challenges of SD are compounded for SR due to the high spatial dimensions the model operates on, and the need to use an image encoder at inference time. Applying SD models for Super-Resolution (SR) introduces additional significant computational challenges that further blow up the computational cost by at least one order of magnitude (compared to standard image generation tasks). Specifically, for practical applications, a SR model is typically applied on large input images (512x512 and higher) that require a “patchification” process whereby the input image is divided into smaller patches, the model is applied to each patch separately and then the results are “stitched” back together to produce the final high resolution output. For instance, a 4x upscaler of an input 512x512 image requires at least 16 times independent model evaluations on 128x128 patches (typically this number is higher as the stitching process requires some overlapping outputs).

[004] The efficiency of stable diffusion (SD) has been widely studied recently, with a primary focus on reducing the denoising steps during inference. Notable efforts include more effective sampling strategies, such as those found in distillation techniques. Prior work have noted that the denoising process for tasks such as super-resolution starts with the low-resolution image rather than random noise, potentially requiring fewer inference steps. Recent studies have shown good performance even with very small step counts, reaching the extreme case of just one step. Nonetheless, the computational cost of a large model prevents deploying such models on resource-limited devices like cell phones.

[005] The applicant has identified the need for an improved method for training a stable diffusion model to perform super-resolution. Summary

[006] In a first approach of the present techniques, there is provided a computer-implemented method for training a stable diffusion machine learning, ML, models, for superresolution, the method comprising: training a ll-Net for a super-resolution task using bidirectional conditioning to condition a low-resolution input image from a training dataset. A super-resolution task may be considered to increase a resolution of a low resolution (LR) image, e.g. to a high resolution (HR) image.

[007] The first approach may thus be a computer-implemented method for training a diffusion machine learning, ML, model for increasing a resolution of an input LR image, wherein the diffusion ML model comprises a LR encoder for generating a LR embedding vector for the input image, a denoising ML model which generates a HR embedding vector and a decoder for generating a HR image from the output of the denoising ML model. The method comprises: obtaining a training dataset comprising at least one high resolution, HR, image; inputting a HR image from the training dataset into a HR encoder to generate a HR embedding vector for the HR image, wherein the HR encoder is separate from the diffusion ML model which is being trained. The LR encoder may be considered to be a first encoder and the HR encoder may be considered to be a second encoder. The method comprises obtaining a LR image from the HR image (e.g. by degrading the HR image using standard techniques); inputting the obtained LR image into the LR encoder to generate a LR embedding vector for the LR image, wherein the LR embedding vector comprises a mean embedding vector and a variance embedding vector for the LR image; obtaining at sequential time steps a plurality of sample latent variables by combining the HR embedding vector with a noise sample taken from a Gaussian distribution defined by the generated mean and variance embedding vectors; predicting, for each sample latent variable, using the denoising ML model, a HR embedding vector; training the denoising ML model by minimising a difference between the predicted HR embedding vector and the generated HR embedding vector; and outputting, from the server to the user device, the trained diffusion ML model which comprises the LR encoder, the trained denoising model, and the decoder for generating a HR image from an output of the trained denoising model.

[008] A diffusion model may also be termed a diffusion probabilistic model or a score-based generative model and is a class of latent variable generative models. A diffusion model typically has three components: a forward process, a reverse process and a sampling process. The goal of a diffusion model is to learn a diffusion process for a given dataset so that the model can generate new elements that are distributed similarly to the original dataset. The present techniques may use a stable diffusion model which comprises a variational autoencoder (VAE), a U-Net and an optional text encoder. The VAE encoder compresses the image from pixel space to a smaller dimensional latent space and Gaussian noise is iteratively applied to the compressed latent representation during forward diffusion. The U-Net block denoises the output from the forward diffusion backwards to obtain a latent representation and the VAE decoder generates the final image by converting the latent representation back into pixel space. In the present techniques, the diffusion model is being used for super resolution, in other words, the diffusion model is being used to generate an output image having a higher resolution than an input image. A higher resolution image may have a magnification of twice, four times or another multiple of the lower resolution image. In other words, the concentration of pixels may be twice, four times or another multiple of concentration of pixels in the lower resolution image Thus, the diffusion model is being used to denoise the input image (also termed a low resolution image) to generate a high resolution output image. The diffusion model may thus comprise a denoising model, for example a U-Net. A U-Net is a convolutional neural network that was developed for image segmentation. Each embedding vector (e.g. the LR embedding vector, the HR embedding vector, the mean embedding vector and the variance embedding vector) is a one dimensional tensors / arrays of the embeddings. The embeddings are the numerical representations of data points within the latent space.

[009] There may also be a check to see whether parameters of the denoising M L model are converging at each sequential time step and when there is no convergence, inputting a further HR image from the training dataset and repeating the steps of obtaining a LR image, inputting the obtained LR image, obtaining a plurality of samples, predicting and training until there is convergence.

[010] The noise sample at each sequential time step may be taken from J\T(z^diag(zf)) where J\T is a Gaussian distribution, z^ are the mean embedding vectors and zf are the variance embedding vectors. By incorporating this type of noise sample which is sampled from the mean and variance embeddings of the LR image into the latent variable, the forward process of the diffusion model is conditioned on the LR image. In other words, the process of adding noise is conditioned on the LR image rather than on a random noise which is sampled in a standard unidirectional conditioning process.

[011] A latent variable is a variable which can only be inferred indirectly through a mathematical model from other observable variables that can be directly observed or measured (in this case the LR embeddings).https: / / en.wiktpedia.orR / wiki / Latgnt and observable variables - cite note-2 The sample latent variable which is input to the denoising model may be a weighted sum of the noise sample and the HR embedding vector. For example, the input latent variable may be defined using the mathematical expression: zt = trfZ / ( + where zt |t e [0,1] is the input latent variable at time t, at,ot are noise parameters which define the noise schedule, e is the noise sample, namely e ~ XCz^diagCz^), and zh is the HR embedding vector. Thus, the weighted sum may use a first weight at (also termed first noise parameter) for the noise sample and a second weight at (also termed second noise parameter) for the HR embedding vector. The noise parameters may define a noise schedule having a signal to noise ratio which decreases with each sequential time step. For example, the log signal-to-noise ratio, At = \og[al / o^] may decrease with each time step t monotonically. In other words, the second weight increases relative to the first weight with each sequential time step during a denoising process of the denoising ML model so that at a final time step, there is effectively no noise sample.

[012] Training the denoising ML model may comprise adjusting parameters of the denoising ML model to minimize the loss or difference between the predicted HR embedding vectors and the generated HR embedding vector. For example, a weighted mean square error loss may be used, e.g. as defined by argmin e Ee<t[&>(At)|[z0(zt,At) - zh||f] where z0 is the denoising ML model, zh is the HR embedding vector, At is the log signal-to-noise ratio, to is a weighting function, e ~ JV(0, / ) a Normal Distribution and zt is the latent variable defined above. The weighting function may be any suitable function, e.g. v parametrization.

[013] The method may comprise training the LR encoder and keeping parameters of the HR encoder frozen. The method may comprise: training a small encoder (LR encoder) and the U-Net and the denoising model jointly. The output trained diffusion model may further comprise the trained LR encoder.

[014] A scheduler may be used to schedule training of the LR encoder jointly (i.e. simultaneously) with the denoising model for at least some of the sequential time steps and schedule training only the denoising model for other time steps. Jointly training the LR encoder and the denoising ML model may comprise adjusting parameters of the denoising ML model and parameters of the LR encoder to minimize the loss between the predicted HR embedding vectors and the generated HR embedding vector. For example, a weighted mean square error loss may be used, e.g. as defined by argmin 9iCf) Eet[a)(At)\\20(atzh + + * <Uf) - zhIli] where the LR encoder is defined by = zf, and ££(xt) = zf and the other parameters are as defined above. Only parameters of the denoising model may be adjusted at these other time steps to minimise a difference between the predicted HR embedding vectors and the generated HR embedding vector

[015] The sequential time steps may be selected from t e [0,1] and there may be approximately 1000 time steps. The sampling may be uniform over time. Alternatively, when joint training is used, the sampling may not be uniform and the sampling rate may be higher for t =1 than when t <1. Joint optimisation may take place only when t =1. In other words, the method may comprise using the scheduler to sample at a higher rate when scheduling training of the LR encoder simultaneously with the denoising model and to sample at a lower rate when scheduling training only of the denoising model. Sampling at a higher rate may comprise decreasing a time gap between successive sequential time steps. Sampling at a lower rate may comprise increasing a time gap between successive sequential time steps, for example, until the time gap converges to a uniform time gap for later time steps.

[016] The method may further comprise training an auto-encoder on high-resolution images from the training dataset, while freezing all layers of the small encoder except a final layer.

[017] The method may further comprise: fine-tuning a decoder; and training a small decoder on top of the jointly-trained encoder and U-Net, to thereby obtain a lightweight single-step model. In other words, the method may further comprise training a decoder before outputting the trained diffusion model and outputting the trained diffusion model with the trained decoder. In other words, the parameters of the decoder may be adjusted to minimise a difference between the HR image which is output from the decoder and the original HR image from the training dataset which was input.

[018] The method may further comprise using scale distillation. Specifically, using scale distillation may include training a teacher model for a lower magnification (e.g. x2) and progressively distilling the teacher to a student model at the target scale factor of interest (e.g. x4). The student model in this example is the denoising model.

[019] The method may further comprise reducing a size of the denoising model before training the denoising model. For example, the number of channels may be reduced, the number of blocks / layers within the model may be reduced and / or the number of parameters within the model may be reduced. Similarly, the method may further comprise reducing a size of the LR encoder, e.g. by reducing channels and / or blocks. For example, there may be 158M parameters compared to 860M parameters, e.g. over a four-fold reduction. The number of channels may be halved. The number of layers within the model may be reduced to one, whereby there is a 1-step inference process when using the model.

[020] The user device which receives the output model, may be any electronic device and may particularly be a constrained-resource device. The reduction in size of the model thus enables its successful deployed on such devices. It is also noted that the HR encoder is only used in training and is not deployed which also reduces the load on the user device. The user device may be a smartphone. A smartphone is an example of a smart device. The user device may be a smart appliance. A smart appliance is another example of a smart device. An example of a smart appliance is a smart television (TV), a smart fridge, a smart oven, a smart vacuum cleaner, a smart robotic device, a smart lawn mower, and so on. More generally, the user device may be a constrained-resource device, but which has the minimum hardware capabilities to use the personalised ML model.

[021] According to another aspect of the present techniques, there is provided a computer-implemented method for generating a high resolution, HR, image from a low resolution, LR, image using the trained diffusion model described above on a user device. The method may comprise receiving and storing an output trained diffusion ML model which has been trained as described above and which comprises the LR encoder, the trained denoising model, and the decoder. The method for generating a HR image may comprise; receiving an LR image (e.g. from a camera on the user device or from a user); inputting the received LR image into the trained diffusion ML model on the user device; generating, using the LR encoder, a LR embedding of the received LR image; inputting the LR embedding into the trained denoising ML model to generate a HR embedding; and generating, using the decoder, a HR image from the generated HR embedding.

[022] According to another aspect of the present techniques, there is provided a user device for generating a high resolution, HR, image from a low resolution, LR, image, the user device comprising at least one processor and memory storing instructions that, when executed by the at least one processor individually or collectively, cause the user device to carry out the method of generating a HR image as described above.

[023] According to another aspect of the present techniques, there is provided a server for training a diffusion machine learning, ML, model for increasing a resolution of an input low-resolution, LR, image,, the server comprising: at least one processor and memory storing instructions that, when executed by the at least one processor individually or collectively, cause the server to carry out the method of training which is described above.

[024] The features described above with respect to the first approach apply equally to this approach and therefore, for the sake of conciseness, are not repeated.

[025] In a related approach of the present techniques, there is provided a computer-readable storage medium comprising instructions which, when executed by at least one processor, causes the processor to carry out any of the methods described herein.

[026] In the cases where the present techniques are implemented or executed on a device comprising multiple processors, the present techniques may be implemented by one or more of the multiple processors. That is, the present techniques may be implemented by or executed by the processors individually or collectively.

[027] As will be appreciated by one skilled in the art, the present techniques may be embodied as a system, method or computer program product. Accordingly, present techniques may take the form of an entirely hardware embodiment, an entirely software embodiment, or an embodiment combining software and hardware aspects.

[028] Furthermore, the present techniques may take the form of a computer program product embodied in a computer readable medium having computer readable program code embodied thereon. The computer readable medium may be a computer readable signal medium or a computer readable storage medium. A computer readable medium may be, for example, but is not limited to, an electronic, magnetic, optical, electromagnetic, infrared, or semiconductor system, apparatus, or device, or any suitable combination of the foregoing.

[029] Computer program code for carrying out operations of the present techniques may be written in any combination of one or more programming languages, including object oriented programming languages and conventional procedural programming languages. Code components may be embodied as procedures, methods or the like, and may comprise subcomponents which may take the form of instructions or sequences of instructions at any of the levels of abstraction, from the direct machine instructions of a native instruction set to high-level compiled or interpreted language constructs.

[030] Embodiments of the present techniques also provide a non-transitory data carrier carrying code which, when implemented on a processor, causes the processor to carry out any of the methods described herein.

[031] The techniques further provide processor control code to implement the abovedescribed methods, for example on a general purpose computer system or on a digital signal processor (DSP). The techniques also provide a carrier carrying processor control code to, when running, implement any of the above methods, in particular on a non-transitory data carrier. The code may be provided on a carrier such as a disk, a microprocessor, CD- or DVD-ROM, programmed memory such as non-volatile memory (e.g. Flash) or read-only memory (firmware), or on a data carrier such as an optical or electrical signal carrier. Code (and / or data) to implement embodiments of the techniques described herein may comprise source, object or executable code in a conventional programming language (interpreted or compiled) such as Python, C, or assembly code, code for setting up or controlling an ASIC (Application Specific Integrated Circuit) or FPGA (Field Programmable Gate Array), or code for a hardware description language such as Verilog (RTM) or VHDL (Very high speed integrated circuit Hardware Description Language). As the skilled person will appreciate, such code and / or data may be distributed between a plurality of coupled components in communication with one another. The techniques may comprise a controller which includes a microprocessor, working memory and program memory coupled to one or more of the components of the system.

[032] It will also be clear to one of skill in the art that all or part of a logical method according to embodiments of the present techniques may suitably be embodied in a logic apparatus comprising logic elements to perform the steps of the above-described methods, and that such logic elements may comprise components such as logic gates in, for example a programmable logic array or application-specific integrated circuit. Such a logic arrangement may further be embodied in enabling elements for temporarily or permanently establishing logic structures in such an array or circuit using, for example, a virtual hardware descriptor language, which may be stored and transmitted using fixed or transmittable carrier media.

[033] In an embodiment, the present techniques may be realised in the form of a data carrier having functional data thereon, said functional data comprising functional computer data structures to, when loaded into a computer system or network and operated upon thereby, enable said computer system to perform all the steps of the above-described method.

[034] The method described above may be wholly or partly performed on an apparatus, i.e. an electronic device, using a machine learning or artificial intelligence model. The model may be processed by an artificial intelligence-dedicated processor designed in a hardware structure specified for artificial intelligence model processing. The artificial intelligence model may be obtained by training. Here, "obtained by training" means that a predefined operation rule or artificial intelligence model configured to perform a desired feature (or purpose) is obtained by training a basic artificial intelligence model with multiple pieces of training data by a training algorithm. The artificial intelligence model may include a plurality of neural network layers. Each of the plurality of neural network layers includes a plurality of weight values and performs neural network computation by computation between a result of computation by a previous layer and the plurality of weight values.

[035] As mentioned above, the present techniques may be implemented using an Al model. A function associated with Al may be performed through the non-volatile memory, the volatile memory, and the processor. The processor may include one or a plurality of processors. At this time, one or a plurality of processors may be a general purpose processor, such as a central processing unit (CPU), an application processor (AP), or the like, a graphics-only processing unit such as a graphics processing unit (GPU), a visual processing unit (VPU), and / or an Al-dedicated processor such as a neural processing unit (NPU). The one or a plurality of processors control the processing of the input data in accordance with a predefined operating rule or artificial intelligence (Al) model stored in the non-volatile memory and the volatile memory. The predefined operating rule or artificial intelligence model is provided through training or learning. Here, being provided through learning means that, by applying a learning algorithm to a plurality of learning data, a predefined operating rule or Al model of a desired characteristic is made. The learning may be performed in a device itself in which Al according to an embodiment is performed, and / o may be implemented through a separate server / system.

[036] The Al model may consist of a plurality of neural network layers. Each layer has a plurality of weight values, and performs a layer operation through calculation of a previous layer and an operation of a plurality of weights. Examples of neural networks include, but are not limited to, convolutional neural network (CNN), deep neural network (DNN), recurrent neural network (RNN), restricted Boltzmann Machine (RBM), deep belief network (DBN), bidirectional recurrent deep neural network (BRDNN), generative adversarial networks (GAN), and deep Q-networks.

[037] The learning algorithm is a method for training a predetermined target device (for example, a robot) using a plurality of learning data to cause, allow, or control the target device to make a determination or prediction. Examples of learning algorithms include, but are not limited to, supervised learning, unsupervised learning, semi-supervised learning, or reinforcement learning. Brief description of the drawings

[038] Implementations of the present techniques will now be described, by way of example only, with reference to the accompanying drawings, in which:

[039] Figure 1 is a schematic diagram showing an overview of bidirectional conditioning;

[040] Figure 2 is a flowchart showing the training steps of the present techniques;

[041] Figure 3 is a flowchart showing the training steps for jointly training an encoder with a denoising model;

[042] Figure 4 is a schematic block diagram for a system for implementing the methods of Figures 2 and 3;

[043] Figure 5 is a flowchart showing the inference steps using a model trained using present techniques;

[044] Figure 6a plots the Frechet inception distance (FID) against number of denoising steps for four different models;

[045] Figure 6b plots the learned perceptual image patch similarity (LPIPS) against number of denoising steps for the same four models as Figure 6a;

[046] Figure 7a plots the Frechet inception distance (FID) against number of denoising steps for three different models; and

[047] Figure 7b plots the learned perceptual image patch similarity (LPIPS) against number of denoising steps for the same four models as Figure 7a. Detailed description of the drawings

[048] Broadly speaking, embodiments of the present techniques provide a method for training a stable diffusion model to perform super-resolution of low or lower resolution images, which can be deployed on resource-constrained devices such as smartphones.

[049] The present techniques introduce a novel approach for training efficient stable diffusion models for super-resolution. While prior work often focused on reducing the number of inference steps, the resulting models still have around 1 billion parameters, which is prohibitive for deployment on constrained devices like phones. The present techniques target a smaller U-Net with 158 million parameters and an autoencoder with 11 million parameters. One technical difference relative to prior techniques is the proposal of bidirectional conditioning for conditioning on low-resolution input. Instead of sampling from a normal Gaussian, sampling is from a Gaussian whose mean and variance are obtained from the low-resolution input. It is shown that bidirectional conditioning allows us to effectively train the encoder for the superresolution task and subsequently replace it with a lighter encoder without compromising performance. The final model, with a large number (millions) of parameters, requires a small number of FLOPS for a single step and achieves competitive results compared to the large model with 1 billion parameters that requires a large number of FLOPS.

[050] Specifically, the present techniques concern Stable Diffusion (SD) for on-device image 5 Super-Resolution (SR). SR approaches based on SD have shown to produce high-quality visual outputs, yet they are entirely unsuitable for on-device deployment due to (1) requiring several inference steps, (2) being very big in terms of model size (1B params), and (3) inducing high computational complexity (>2000 GFLOPs). The present techniques address all aforementioned three limitations, developing Edge-SD-SR that requires 1 inference step, uses 10 a lightweight model of total of 170M parameters for both UNET and Encoder / Decoder and has complexity of fewer GFLOPs (only around 142 GFLOPS). The main technical novelty is the proposal of bidirectional conditioning for conditioning on low-resolution input. Instead of sampling from a normal Gaussian, sampling is from a Gaussian whose mean and variance are obtained from the low-resolution input. It is shown that bidirectional conditioning allows us 15 to effectively train a lightweight UNET and encoder without compromising visual quality. It is shown that Edge-SD-SR can 4x upscale an 128 x 128 patch on-device in a short time (38 msecs) and 4x upscale a whole 512x512 image to 2,048x 2,048 (requiring 25 model evaulations) on-device in just ~ 1 sec on a Samsung S24 smartphone. On top of this, it is shown that Edge-SD-SR matches or even outperforms all state-of-the-art SR approaches on 20 the most established SR benchmarks (e.g. DIV2K, RealSR, DRealSR).

[051] Although there are a few attempts to reduce the number of inference steps, due to the complexity associated with the SR task, it is desired to reduce this to just a single step with minimal degradation in quality. To address the aforementioned challenges, the present 25 techniques both reduce model size and complexity of the whole architecture (U-Net and encoder / decoder) and use just a single step during SR model inference, even while producing SR results of high quality. This is different from general-purpose methods for reducing inference steps for diffusion that have been shown to be effective for large models, for example as described in “DPM-Solver +++: Fast solver for guided sampling of diffusion probabilistic 3D models” by Lu et al published in ArXic 2211.01095 (2023). It is also different from methods that focus on reducing model size via architectural changes but still require multiple inference steps, for example as described in “SnapFusion: Text-to-image diffusion model on mobile devices within two seconds” by Li et al published in Advances on NIPS (2023). 35

[052] The present techniques aim to reduce the size of all three components of an SD-SR v1.5 model, i.e., the encoder, U-Net, and decoder, while minimizing the drop in performance. A large model typically includes 65M, 860M, and 35M parameters for each component, respectively, requiring a large number of FLOPS to process a 128 x 128 patch using a single ddim step. The present techniques consider a model that includes less than 200M parameters overall. To achieve this, a small U-Net is obtained with 158M parameters by reducing the number of channels by a factor of four. Unfortunately, reducing the U-Net capacity to this level results in a significant drop in performance. Moreover, the autoencoder latency is relatively high compared to the U-Net. To address the performance drop and reduce the encoder and decoder latency, bidirectional conditioning is proposed that utilizes the low-resolution (LR) input through a more intuitive and efficient mechanism. In other words, at the heart of the method, there is a new training method that tailors diffusion to the task of SR by bidirectional conditioning on the low-resolution input.

[053] A common practice in utilizing diffusion models for super-resolution tasks is to provide the LR input as an additional input via concatenation to the denoising backward process. However, this process does not condition the model explicitly on the LR image. In other words, the forward process is LR-agnostic. To explicitly condition the model on the LR image, conditioning the sampling Gaussian distribution in the forward process on the LR image is proposed. That is, instead of sampling from a normal distribution, the mean and variance of the Gaussian is conditioned on the LR image to generate the noisy version of the data. In the context of SD models, sampling is used to form a Gaussian where its mean and variance are determined by the encoder, mixing with the high-resolution (HR) embeddings through the scheduler to obtain the noisy input for the U-Net. It is empirically shown that the proposed conditioning improves the performance of the small model.

[054] A bottleneck of utilizing text-to-image SD models for super-resolution tasks is that the encoder suffers from the domain shift between the HR and LR images, making it prone to suboptimal embeddings for the LR images. A common solution to address this issue is training an adapter along with the U-Net, which adds more computational overhead to the encoding process. In contrast, the present bidirectional conditioning allows training the lightweight encoder jointly with the U-Net, significantly improving performance. As detailed below subsequently, the trained encoder and U-Net are combined with a lightweight decoder, without compromising visual quality. It is empirically shown that joint training of the encoder and U-Net is feasible due to our conditioning method by repeating the same procedure for the commonly used concatenation approach. It is observed that the gain is significantly in favor of our bidirectional conditioning method. Moreover, it is shown that by jointly training the encoder and U-Net, it is possible to replace the encoder with a smaller model without compromising performance.

[055] To further improve performance, the auto-encoder latent space channels is increased to 8. As a result, the input / output dimensions of the U-Net will be the same as the concatenation, however, the enlarged latent space provides more flexibility that the present bidirectional conditioning approach benefits from that. The 8-channel auto-encoder is trained on the HR images of the present training set while keeping the encoder frozen except for the last layer.

[056] Finally, following a decoder fine-tuning process, a small decoder is successfully trained on top of the jointly trained small encoder and U-Net to obtain a lightweight single-step model. The final lightweight model achieves competitive results with the large model while having significantly fewer parameters and consuming fewer FLOPS.

[057] To summarize, the present techniques make at least the following contributions: • Bidirectional conditioning to effectively train a small U-Net for SD-SR task that is shown to be effective when using a single inference step. • A mechanism for the joint training of the encoder and U-Net for an improved and more efficient encoder. The lightweight encoder and U-Net are then combined with a lightweight decoder giving rise to fully efficient SD-SR architecture. • A final small model (coined Edge-SD-SR) achieves competitive results with state-of-the-art diffusion based SR models (such as that described in “Exploiting diffusion prior for real-world image super-resolution” by Wang et al published in ArXiv 2305.07015 (2023)) on the most established SR benchmarks (e.g. DIV2K, RealSR, DReal SR) and the final small model includes only 169M parameters, and has fewer FLOPS (just 149) and can 4x upscale a whole 512x512 image to 2,048x2,048 on-device in just 1.1 seconds.

[058] Method

[059] In this section, the conditioning method of stable diffiusion (SD) for the super-resolution task is first discussed. Then the present bidirectional conditioning approach for superresolution using stable diffusion framework is described. Finally, the methodology for fine tuning the encoder and decoder is discussed.

[060] Figure 1 is a schematic diagram showing an overview of bidirectional conditioning. A high resolution (HR) image 100 is converted to a low resolution (LR) image 102. The HR image 100 is input to a first, HR encoder 110 (labelled Ehr) and the LR image 102 is input to a second, LR encoder (labelled E) which is different from the first encoder. The embeddings from the LR encoder are used to condition the Gaussian mean and variance. A denoising model 122 (which is also termed a diffusion model and the terms may be used interchangeably) takes zt as an input and is trained jointly with the second, LR encoder. During training, the parameters of the first, HR encoder 110 are fixed. The diffusion model may use a U-Net architecture. The U-Net architecture is typically based on a fully convolutional neural network whose architecture has been modified and extended to work with fewer training images and to yield more precise segmentation.

[061] Figure 2 is a flow chart setting out the steps of the joint training of the LR encoder and the diffusion model. A dataset 200 is obtained in the form of pairs of low-resolution (LR) and high-resolution (HR) images (xh,xt) ~ p(xh,xi), where xft is the HR image andx( is the paired LR image which as explained below may be generated during the training. In other words, the dataset 200 may just comprise a HR image. The objective of a super-resolution (SR) framework is to estimate the probability distribution of p(xh|x{). During the training, x{ is obtained by applying synthetic degradation on xh. As shown in Figure 1, the model comprises two encoders, £ir,£hr for the HR and LR images accordingly. Before the training process, £hr is initialised with the SDv1.5 text-to-image encoder and then frozen. £tr is initialised similarly when using a big encoder for the LR image, and initialized from scratch in the case of the small encoder.

[062] The training process begins at step S202 by inputting a sample HR image xh from the dataset. Two separate processes are then performed on the HR image xh. At step S204, a LR image x, is obtained from the input sample HR image xh. The degraded LR image can be obtained using any suitable technique, including down sampling, adding noise or applying synthetic degradation. At step S206, the LR image x( is input to the LR encoder £tr to obtain the LR embedding vector z(. In a separate step, which may be carried out simultaneously or sequentially with step S206, at step S216, the original HR image xh is input to the HR encoder £hr to obtain the HR embedding vector zh. The HR encoder is fixed during the training process. The LR embedding vector z( comprises the mean embeddings z^ and variance embeddings zf where z( =< zf.zf >.

[063] Bidirectional conditioning: The forward process of the diffusion model, q(z|zh), is defined as a Markovian Gaussian process defined as the following: gCz^z^^,) = J^(zt;atzh + atzf, CTt(zf)), z = {zt|t e [0,1]} (1) where z is the latent variable of the diffusion model, zt is the value of z at timestep t, z^ =< zf,zf >(i.e. z( the LR embedding vector is defined by its mean embeddings zf and variance embeddings zf) and zh,zf, zf are of the same dimensions, where zh is the HR embedding vector. at, at define the noise schedule such that the log signal-to-noise ratio, At = log[®t2M2]. decreases with t monotonically. For the noising process, it will be appreciated that at time t = 0, there will be no noise and at time t = 1, the input is all noise. There are typically a finite number of steps between 0 and 1, e.g. 1000.

[064] During training, the model learns to reverse this diffusion process progressively, i.e. to estimate p(zt_1\zt), to generate the super-resolved image without needing z( as an additional input given that it is now included in zt. It is possible to estimate the backward denoising process p(zh |zt, zt) using a denoising model in the form of a neural network zg via a weighted mean square error loss: argmin e t[<^(At)||z0(zt, At) - zh|||] (2) over uniformly sampled times t e [0,1] where ze is the neural network and ze(zt,At) may also be termed zpred, zh is the HR embedding vector, A, is the log signal-to-noise ratio, a> is a weighting function, e ~N(0,I) a Normal Distribution and zt is the latent variable of the diffusion model defined by: zt = atzh + orf + otz^ * e = atzh + ote, and e ~ W(zf, diag(zf )). (3) where zh is the HR embedding vector, zf are the mean embeddings and zf are the variance embeddings zf for the the LR embedding vector z( and at, at are parameters which define the noise schedule. For the denoising process, it will be appreciated that at time t = 0, the input is all noise and at time t = 1, there will be no noise. The weighting function may be any suitable weighting function, such as v parameterization described in “Progressive distillation for fast sampling of diffusion models” by Salimans et al published in International conference on Learning Representations (2022), where: at = 1 1--T °?

[065] In other words, as shown in Figure 2, there is a sampling process which begins at step S208 in which a plurality of samples of the input latent variable are obtained, one for each different time step. Each sample comprises a sample derived from the LR image which is obtained from a Gaussian distribution (also termed a Normal distribution) which is defined by the mean and variance embeddings obtained in step S206, i.e. e ~ JV(zf,diag(zf))). In other words, the mean embedding is used as the mean of the distribution and the diagonal of the variance embedding is used as the variance of the distribution. The denoising model zfl is then trained at step S222 as explained above using a sampling process which includes the sample derived from the LR image and the HR embedding vector which are in a weighted sum and the weights of the sum define the noise schedule. As shown at step S224, there is a check to see whether the model z0 is converging and the error between the predicted output and the expected output is optimised. If the convergence is satisfied, the model is output at step S226. Otherwise, the method loops back to the start to select another image to repeat the whole process.

[066] A direct benefit of the proposed formulation is that the resulting denoising model z0 predicts the HR image, at inference time, starting directly from the LR image instead of noise. Precisely, during inference, we start from za ~ JV^zf, diag(zf)) and call the denoising function progressively. Similar to the inference process from a trained model diffusion model, an iterative approach is followed for sampling through a series of sequential calls, steps, of z0. The process starts by sampling from e ~ JVXzf, diag(zf)) where the quality of the generated image improves monotonically with the number of steps. The widely used denoising diffusion implicit model (DDIM) sampler is used and an example of the DDIM sampler is described in “Denoising diffusion implicit models” by Song et al published in International Conference on Learning Representations in 2021. DDIM is a method to take any model trained on denoising diffusion probalistic model (DDPM) loss, and use it to sample with some steps skipped, sacrificing an adjustable amount of quality.

[067] Comparison with unidirectional conditioning: The main difference between the present proposed method and the commonly used unidirectional conditioning approach via concatenation lies in the forward process. The common approach uses the following Markovian Gaussian process: ?(zt|z / i) = 77(zt; atzh, Cfl), z = {zt|t e [0,1]} (4) where zh = £(xh) and z( = S(x() are the embeddings of the HR and LR images, respectively, given that stable diffusion operates on the latent space of an encoder £. Note that the above forwarding noising process is the same as for the text-to-image case. However, compared to the text-to-image case, diffusion models for super-resolution do not condition on an input caption. Instead, they incorporate a mechanism for conditioning on the LR image. Here, zt, and consequently the forward process, does not depend on the LR image.

[068] To condition the output on the LR image, the LR image is given to the backward denoising function as an additional input. This is typically done through concatenation with the noisy input. In other words, in unidirectional conditioning, only the backward denoising process p(zt_1|zt) is conditioned on the LR image through the following objective function: argmin e Ee7t[6)(At)||zfl(zt,zbAt) -zh||^] (5) over uniformly sampled times t e [0,1] where z0 is the neural network, zh is the HR embedding vector, z( is the LR embedding vector, At is the log signal-to-noise ratio, to is a weighting function, 0 are the parameters of z0 and zr is the input to the neural network and is defined by: zt = atzh + ate; 6 -7 / (0,1) Starting from zt -7 / (0,1), the inference process involves a series of sequential calls wherein the denoising function is called progressively. The denoising function takes the LR image z( as an additional input.

[069] The mechanism defined by equations (4) and (5) is termed unidirectional because the forward process is not conditioned on the LR image. During the inference of the proposed techniques (termed SD-SR), sampling is from e ~ JV(zf, (zf)) and the denoising function is called recursively, not using zt as an additional input. By contrast, LDM-SR samples from 6 -7 / (0,1) and uses zt as an additional input in each recursive call of the denoising function. In essence, SD-SR does not estimate p(zh\zt) explicitly. The denoising process is guided by the LR image to generate the relevant output. In the present formulation, however, the objective function is explicitly optimized for the SR task.

[070] Impact of noise variance conditioning: In the proposed bidirectional conditioning, the forward process is conditioned on the LR image, while taking into account both the mean and variance as explicitly outlined in equations (1) and (3). It is noted that “Inversion by direct iteration: An alternative to denoising diffusion for image restoration” by Delbracio et al published in ArXiv 2303.11435 (2023) proposes to learn a direct transition between LR and HR image, while assuming no knowledge of the forward degradation process. The formulation in (3) is inline with this paper as it also learns a direct transition between the LR and HR image. However, in the proposed techniques, the degradation process is kept the same as in diffusion models via conditioning the sampling Gaussian on the LR input. For comparison, a direct translation of the formulation by Delbracio to the stable diffusion framework, results in the following forward process: zt = atzh + <7tz^ + ot6 = atzh + ote, and e ~ 3 / (0,1) e ~ 3 / (1^,V). (6)

[071] This degradation process is restricted compared to the proposed techniques because only the mean of the sampling Gaussian distribution is conditioned on the LR image. It is proposed that taking the LR image into account for the noise variance is of paramount importance for the SR task to recover details in the resulting HR image. The impact of using equations (3) or (6) is studied below. Significant superiority in the formulation of (3) is found.

[072] Encoder Training: The known SR methods based on stable diffusion use the same encoder £ for the HR and LR image, pretrained on the HR images. To cope with the domain shift between the LR and HR images, an adaptor is widely used trained jointly with the U-Net, resulting to computational overhead.

[073] As described above, the present techniques use two different encoders for the HR embeddings and the LR embeddings and the HR encoder is frozen. The present techniques also propose training the LR encoder jointly with the diffusion model (U-Net) for the SR task, which allows not only computational overhead to be avoided but also reduces the encoder size drastically. The LR encoder can be trained by adding its parameters to the optimization using the same equation 2: argmin Eejt[w(At)||za(atzh + ^£^) + ot£^(xt) * e,At) - zh||^] (7) where £%(Xl) = z*, and £^(xt) = zf.

[074] Preventing the degenerate solution. The encoder is time-agnostic and as a result, the vanilla training of the objective function in Eq. 7 results in a degenerate solution for the average t. To prevent this, a stop gradient technique is performed for t <1 for the encoder which is illustrated in Figure 3. As shown at step S300, we obtain the start values of the mean and variance LR embeddings. In other words, we obtain z± ~ J\f(z^,diag(zf)). Similarly at step S302, which is shown being performed in parallel but can be done sequentially before or after step S300, we obtain the start value of the HR embedding, namely zh. At the first iteration as shown at step S304, there is a determination of whether t = 1.

[075] When t = 1, optimisation is for both 0, ¢, as shown at step S306. In other words, the parameters 0 of the denoising model ze and the parameters ¢ of the encoder model are jointly optimised. By contrast, as shown at step S308, optimisation is only for ¢ when t <1. Regardless of which optimisation is performed, there is a determination as to whether there are more time steps to be sampled at step S310. Since t is uniformly sampled from a large number of discrete steps in practice, 1000 steps, the solution above is prone to reduce to Eq. 2 where only the denoising model is optimized. To address this issue, instead of sampling uniformly over t e [0,1], a scheduler is used so that the sampling is more frequent for t = 1 in the beginning, and the t = 1 sampling rate is gradually dropped to uniform sampling in the later training iterations. Thus, if there are more sampling steps scheduled as determined at step S310, the method loops back to step S304 to determine whether t = 1.

[076] If there are no more sampling steps scheduled, there is a check to see whether there is convergence of the model at step S312 and if so, the optimised denoising model and encoder are output at step S314. If there is not convergence of the denoising model, the process returns to the initial sampling step and a different input may be selected to begin the training process again.

[077] It will be appreciated that the joint training of the encoder described above, can be potentially applied for the SD-SR framework. However, it is significantly more effective when applied with the present proposed bidirectional conditioning method.

[078] As shown in Figure 3, an optional final step is included. At step S316, the final step of the pipeline is to train the decoder on top of the efficient 1-step model (i.e. output trained encoder and denoising model). The fine tuning may be done using any suitable technique, e.g. using the same VQAE loss used to train the original auto-encoder as described in “High-resolution image synthesis with latent diffusion models” by Rombach et al published in IEEE Conference on Computer Vision and Pattern Recognition. It will be appreciated that this optional step may also be used in the method without the dual training of the encoder and denoising model.

[079] Scale distillation

[080] The proposed approach(es) can be readily combined with other techniques specifically targeting speeding up inference of the SD model, such as faster samplers (e.g. as described in “Denoising diffusion implicit models” by Song et al published in International Conference on Learning Representations in 2021, progressive distillation as described in “Progressive distillation for fast sampling of diffusion models” by Salimans et al published in International Conference on Learning Representations in 2022 or scale distillation as described in “Fast super-resolution with stable diffuion via scale distribution” by Noroozi et al published on arXiv 2401.17258 in 2024. As shown below, scale distillation is used because it was specifically designed for the SR task and demonstrated superior performance. Specifically, the method included training a teacher model for a lower magnification (e.g. x2) and progressively distilling the teacher to a student model at the target scale factor of interest (e.g. x4).

[081] Architecture

[082] Figure 4 is an example system for implementing the methods described above. A first electronic device 400 is used to implement the training method described above. The electronic device 400 comprises: at least one processor 402 coupled to a memory 404. The at least one processor 402 may comprise one or more of: a microprocessor, a microcontroller, and an integrated circuit. The at least one processor 402 may include one or more central processing units (CPUs) and / or one or more graphics processing units (GPUs). The memory 404 may comprise volatile memory, such as random access memory (RAM), for use as temporary memory, and / or non-volatile memory such as Flash, read only memory (ROM), or electrically erasable programmable ROM (EEPROM), for storing data, programs, or instructions, for example. The electronic device may be any suitable system or device, e.g. a server.

[083] The electronic device 400 comprises the full ML model 410 for training on device using the paired images in the dataset 420. The dataset 420 is shown as external to the electronic device 400 but may be stored within memory on the device, with at least one set of weights. There is also an input / output interface for receiving data, e.g. from the dataset 420 or for transmitting information to a user device 450.

[084] The full model 410 comprises an autoencoder 412 which includes the HR encoder which generates the HR embeddings which are used in the training process as described above in relation to Figure 2. The autoencoder 412 is fixed during the training of the rest of the model. The full model 410 also a lightweight model which may be termed an Edge-SD-SR model and which is for deployment to one or more user devices 450. The lightweight model comprises the LR encoder 414 and diffusion model 416 which may be jointly trained as detailed above in Figure 3. There is also a LR decoder 418 which may be fine tuned after the rest of the lightweight model is trained.

[085] The lightweight model, once trained, is deployed on the user device 450. The user device may also comprise the standard components such as processor(s) 452, memory 454 and an I / O interface 458. The user device may also have an image capture device 456, e.g. a camera. The user device which receives the lightweight model, may be a constrained-resource device. The device may be any one of: a smartphone, tablet, laptop, computer or computing device, virtual assistant device, a vehicle, an autonomous vehicle, a robot or robotic device, a robotic assistant, image capture system or device, an augmented reality system or device, a virtual reality system or device, a gaming system, an Internet of Things device, or a smart consumer device (such as a smart fridge, smart vacuum cleaner, smart oven, or smart lawn mower). It will be understood that this is a non-exhaustive and non-limiting list of example devices.

[086] The diffusion model may include a standard LI-Net or may be modified by trimming the expansion ratio in the number of channels and / or using transformer blocks in lower resolutions. Trimming may include reducing the expansion ratio by a half. Additionally, a reduction in the number of base channels (e.g. halving) is also proposed because as shown in the results, this approach yields the biggest gain in terms of parameters and FLOPS. These modifications yield an efficient diffusion model that has only 158M parameters and uses around 40 GFLOPS to process an 128x128 patch (compared to 860M parameters and 222 GFLOPS for an unmodified U-Net such as SD 1.5.

[087] Similarly, the auto-encoder used in SD 1.5 is also heavy with about 83M parameters and 1820 GFLOPS. Therefore, the autoencoder used in the proposed techniques is also reduced in size. For example, the number of channels and residual blocks are halved. One downsampling layer may also be removed from the LR encoder 414 so that it directly scales down the input image of size 128x128 by 4. By contrast, the HR encoder scales down HR images of size 512x512 by 8. Importantly, given the proposed decoupling of the HR and LR encoders, the HR encoder can be kept frozen and is only required at training time, while we can deploy the efficient LR encoder 414 at inference time. These design choices result in a lightweight auto-encoder with 14M parameters that only require 102 GFLOPs to process a 128x128 patch.

[088] When the Edge-SD-SR model is deployed from the electronic device 400 to the user device 450 standard quantization techniques may also be used. For example W8A16 quantization which uses vanilla weights - int8 and activations - int16 may be used. Such a deployment means that the deployed model can 4x upscale a 128x128 patch on S24 NPU (QNN 2.19) in approximately 38 msecs (the time at the encoder is 7.465 msecs, the diffusion model 8.340 msecs and the decoder is 22.453 msecs). Moreover, to 4x upscale a whole 512 x 512 image to 2,048 x 2,048 (25 model evaluations in total) this takes approximately 1.1 secs (including the tiling operations). Notably, there is no observable quantitative or qualitative degradation by the quantization process. Improved latency could be achieved by further quantizing the activations.

[089] Figure 5 shows the steps which are carried out at inference time on the user device to which the lightweight model has been deployed. In a first step, the user device captures an image at step S500. The captured image may be a low quality image of a distant object which has been captured using an optical zoom function of the camera on the user device. Optionally, the captured image may be divided into a plurality of patches at step S402. For example, the original captured image may be a 512x512 image and each patch may have a size of 128 x128 pixels. There may be patches which overlap to assist with generating the final output. At step S504, the input image (or patches when used) are input to the lightweight LR encoder to generate the LR embeddings which include the mean and variance embeddings as described above.

[090] The LR embeddings are then input at step S506 to the denoising model which predicts and outputs the HR embeddings at step S508. Precisely, during inference, we start from zt ~ JVXzf, diag(zf)) and call the denoising function progressively. An iterative approach is followed for sampling through a series of sequential calls, steps, of z0. The quality of the generated image improves monotonically with the number of steps. The HR embeddings are then decoded by the decoder at step S510 to generate the HR image (or HR patch where used). A final optional step S512 is to create the final output HR image using the patches. The final HR image is then output at step S514.

[091] Experiments

[092] In this section, the proposed method is extensively evaluated. The training datasets use a combination of DIV2K described in “Ntire 2017 challenge on single image superresolution: Dataset and study” by Agustsson et al published in IEEE Conference on Computer Vision and Pattern Recognition in 2017, DIV8K described in “Div8k: Diverse 8k resolution image dataset” by Gu etal published in IEEE Conference on Computer Vision in 2019, Flickr2k described in “Ntire 2017 challenge on single image super-resolution: Methods and Results” by Timofte etal published in IEEE Conference on Computer Vision and Pattern Recognition in 5 2017, OST described in “Recovering realistic texture in image super-resolution by deep spatial feature recognition” by Wang et al published in in IEEE Conference on Computer Vision and Pattern Recognition in 2018 and a subset of 10K images from the FFHQ training set described in “A style-based generator architecture for generative adversial networks” by Karras et al published in in IEEE Conference on Computer Vision and Pattern Recognition in 2019. The 10 Real-ESRGAN degradation pipeline described in “Training real-world blind super-resolution with pure synthetic data” by Wang et al published in in IEEE Conference on Computer Vision in 2021 is used to generate synthetic LR-HR pairs.

[093] The present model is evaluated on both synthetic and real datasets. 3K LR-HR (128 -> 15 512) pairs synthesized from the DIV2K validation set using the Real-ESRGAN degradation pipeline are used as the synthetic dataset. Results are also reported on the standard DIV2K validation split with bicubic degradations for completeness. For the real dataset, 128 x 128 center crops are used from the RealSR, DRealSR and DPED-iphone datasets. The RealSR dataset is described in “Real-world super resolution via kernel estimateion and noise injection” 20 by Ji et al published in IEEE Conference on Computer Vision and Pattern Recognition in 2020, the DRealSR dataset in “Gans trained by a two time-scale update rule converge to a local nash equilibrium” by Heusel et al published in Advances on Neural Information Processing Systems in 2017 and the DPED-iphone dataset in “Dslr-quality photos on mobile devices with deep convolutional networks” by Ignatov et al publihsed in IEEE Conference on Computer 25 Vision in 2017.

[094] Baselines: As the main contribution of the present techniques targets improving the inference process of diffusion-based super-resolution, the main points of comparison are diffusion-based SR models, including the recent StableSR model, Yonos-SR and ResShift 3D and the original LDM model. The StableSR model is described in “Exploiting diffusion prior for real-world image super-resolution” by Wang et al published in arXiv 2305.07015 in 2023, the Yonos-SR model is described in “You only need one step: Fast super-resolution with stable diffusion via scale distillation” by Noroozi et al published in arXiv 2401.17258, the ResShift model is described in “Efficient diffusion model for image super-resolution by residual shifting” 35 by Yue et al published in Advances on Neural Information Processing Systems (2023) and the original LDM model is described in “High-resolution image synthesis with latent diffusion models” by Rombach et al published in IEEE Conference on Computer Vision and Pattern Recognition in 2022.

[095] Evaluation metrics: Evaluation involves using various perceptual and image quality metrics, including learned perceptual image patch similarity (LPIPS), Frechet inception distance (FID) (where applicable), as well as the no-reference image quality metric, MUSIQ which is described in “Musiq: Multi-scale image quality transformer” by Ke et al published in IEEE Conference on Computer Vision in 2021. Standard peak signal-to-noise ratio (PSNR) and SSIM metrics are reported, for reference.

[096] Figure 6a plots the Frechet inception distance (FID) against number of denoising steps for four different models and Figure 6b plots the learned perceptual image patch similarity (LPIPS) against number of denoising steps for the same four models as Figure 6a. These two metrics are chosen because of their alignment with human evaluation of the resulting HR images. By varying the number of denoising steps, the impact of each component is subject to a more indepth assessment. Specifically we consider 1 to 8 steps at inference time. For this study, the set of 3K LR-HR pairs synthesized from the DIV2K validation set is used to extensively study the impact of our contributions on a controlled large test set. All the results are obtained using the original SD 1.5 big decoder.

[097] The impact of each component in the proposed techniques is analysed via pairwise comparisons. Each solid line corresponds to results obtained with one of our proposed componetns and the corresponding dashed line represents the corresponding baseline without this component. Thus, the four models which are shown in Figures 6a and 6b are: solid line with circles - a large encoder model having four channels trained using the proposed bidirectional conditioning; dashed line with circles - a large encoder model having four channels trained using standard unidirectional conditioning; solid line with squares - a large encoder model having four channels trained using the proposed bidirectional conditioning and joint encoder training; dashed line with squares - a large encoder model having four channels trained using standard unidirectional conditioning and joint encoder training.

[098] Bidirectional conditioning. Figures 6a and 6b evaluate the impact of the proposed bidirectional conditioning strategy while training a more efficient U-Net (diffusion model). To this end, two variants of the small U-Net are trained while keeping the auto-encoder frozen. Specifically, one U-Net is trained with the proposed bidirectional conditioning (solid line with circles) and another U-Net is trained with the standard unidirectional conditioning (dashed line with circles). The standard unidirectional conditioning means conditioning on the LR image via concatenation as explained above. All other training details are identical. The results depicted in the solid line with circles speak decisively in favour of the proposed bidirectional conditioning, which outperforms the widely used unidirectional conditioning (dashed line with circles) with significant margins for various numbers of denoising steps.

[099] Encoder training. Figures 6A and 6B also show the impact of the joint training of the encoder with the U-Net as described for example in Figure 3. To evaluate the effectiveness of the encoder training in conjunction with the proposed bidirectional conditioning, the encoder is trained jointly with the U-Net using both conditioning strategies (i.e. unidirectional vs bidirectional). When training the encoder of the LR image, the HR image encoder is kept frozen as described above. Comparing the solid line with squares to its dashed counterpart, we can see that although encoding training is possible for both conditioning approaches, it is more effective in the case of bidirectional conditioning. Importantly, the effectiveness of encoder training with bidirectional conditioning is especially striking for the challenging scenario of interest, where we target using only 1 denoising step at inference time. It is possible that bidirectional conditioning provides a more stable training signal (i.e. less corrupted by noise and closer to the target HR), thereby yielding more effective joint training. In light of these results, joint encoder-UNet training using bidirectional conditioning is adopted for all the following experiments.

[100] Encoder size and channels. Figures 7A and 7B show the impact of the encoder size and channels. Figure 7a plots the Frechet inception distance (FID) against number of denoising steps for three different models and Figure 7b plots the learned perceptual image patch similarity (LPIPS) against number of denoising steps for the same four models as Figure 7a. The three models which are shown in Figures 7a and 7b are: solid line with circles - a large encoder model having four channels trained using the proposed bidirectional conditioning and joint encoder training; solid line with squares - a large encoder model having eight channels trained using the same training; solid line with triangles - a small encoder model having eight channels trained using the same training.

[101] As it has been shown in the text-to-image task described in “Enhancing image generation models using photogenic needles in a haystack” by Dai et al published in arXiv 2309.15807 in 2023, increasing the encoder size from 4 to 8 results in improved performance in the SD model. The same procedure is followed for the SR task and the number of channels in the last layer of the encoder and first layer of the U-Net is modified from 4 to 8. As shown in the Figures, training the 8-channel encoder jointly with the U-Net for the SR task improves the performance further compared to the same experiment using 4-channel (see circles versus squares plots).

[102] The big encoder, however, includes 65M parameters, consuming many FLOPS to process a 128 x 128 patch. To reduce the encoder size and latency, a small encoder is created by reducing the number of channels and residual blocks to half, resulting in a small encoder with 11M parameters that require only fewer FLOPS to process a 128 x 128 patch. As shown in the triangle plot in the Figures, despite a drastic reduction in the model size and latency, the trained small / lightweight encoder performs competitively with the big encoder.

[103] Impact of noise variance conditioning. As discussed above, the variance of the sampling Gaussian could be set to identity and the formulation is then similar to equation (6). To study the impact of variance when conditioning the sampling Gaussian on the LR image, the sampling variance is set to the identity and the lightweight U-Net with 8 channel big encoder is then retrained. The results are shown in Table 1 below and removing the variance from the conditioning results in a significant drop in performance, thereby emphasising the benefits of the proposed formulation. Method Non-conditioned noise variance Conditioned noise variance FID^ 67.7 59.2 LPIPSl 0.47 0.44 PSNRt 24.4 24.8 SSIMt 0.63 0.60 Table 1 - Impact of noise variance conditioning. Results with 1-step using the big original decoder.

[104] Comparison with the state-of-the art. The performance of the proposed model is evaluated in the standard real image super-resolution setting targeting x 4 scale factor. Table 2 below focuses on comparisons to SD-based models because the main point of contribution is for diffusion based SR models. The results in bold are the best results and the results in italics are the second best results. Datasets Metrics LDM StableSR YONOS ResShift New Model DIV2K Valid RealESRGAN degradations FID^ LPIPSt PSNRt SSIMf MUSIQf 26.47 0.2510 23.32 0.5762 62.27 24.44 0.3114 23.26 0.5726 65.92 21.86 0.2310 24.74 0.6428 70.30 30.45 0.3076 24.62 0.6210 63.58 25.37 0.249 24.10 0.617 69.58 Real SR LPIPS^ MUSIQt 0.3159 58.90 0.3002 65.88 0.2479 69.21 0.3279 59.87 0.278 65.20 DRealSR LPIPS^ MUSIQf 0.3379 53.72 0.3284 58.51 0.2721 66.26 0.3870 54.13 0.292 55.66 DPED-iphone MUSIQt 44.23 50.48 59.45 38.59 60.09 # STEPS^ 200 200 1 4 1 Table 2: Comparison to stable diffusion-based baselines

[105] It is important to note that the comparison shown in Table 2 is not fair because our proposed model operates on a challenging setting with a very small model and only considering results obtained with one denoising step. By contrast, each of the compared methods only targets one of the aspects that our proposed method tackles collectively. All SD-SR comparing methods consume >2000 GFLOPS, whereas our model requires just 142 GFLOPS. Specifically, the original LDM model uses a similarly sized ll-Net (but larger encoder) and requires 200 denoising steps. Nevertheless our proposed model outperforms this model in all datasets and metrics. StableSR uses a much larger model and 200 denoising steps but focusses on encoder training and once again, the proposed model outperforms StableSR in 7 out of 10 comparison points. YONOS-SR uses a much larger model but specifically optimises for inference speed to obtain a model running in only 1 denoising step. This is a challenging comparison point because our method also tackles inference speed but with a very limited computational budget and thus our method falls behind YONOS-SR in several comparison points. ResShift proposes a new conditioning strategy to rely more heavily on the LR image and uses only 4 denoising steps. Notably our method uses a smaller model size (considering the autoencoder size as well), and our method outperforms ResShift in 8 out of 10 comparison points.

[106] Qualitative testing also supports the superior or on-par results even though our model is significantly smaller and faster than state of art baselines. Collectively these results clearly support the effectiveness of the proposed approach to enable SD-based SR models that can readily run on-device.

[107] Finally, Table 3 below shows the impact of adding scale distillation on top of the proposed bidirectional encoding. As shown in the table, scale distillation has some positive impact primarily on reducing FID. The results in the table were obtained using the DIV2K validation set and used joint training of the small encoder with the U-Net followed by small decoder finetuning. Method New method without scale distillation New method with scale distillation FID I 28.64 25.37 LPIPS 1 0.256 0.249 PSNR T 24.05 24.10 SSIM T 0.619 0.617 Table 3: Impact of scale distillation

[108] Related work.

[109] Image Super-resolution: Image super-resolution involves the process of enhancing a low-resolution (LR) image to reconstruct a high-resolution (HR) version of it. Training one such model typically requires LR-HR image pairs, which are difficult to collect in practice. Early super-resolution methods operate under the assumption of a predefined degradation process, e.g. bicubic downsampling, with known parameters. However, different imaging devices and image signal processing methods result in substantially different degradation processes, and such early approaches fail to generalize to real-world scenarios. Blind super-resolution focuses instead on a more realistic scenario, where the image degradation parameters used to create the low-res image are not assumed known, and the aim is to mimic real-world scenarios. Currently, the most widely adopted approach is the so-called explicit degradation modeling approach, where a degradation pipeline is explicitly constructed and used at training time as an image augmentation to generate artificial LR-HR pairs on the fly, with the expectation being that they resemble real-world LR-HR pairs. The RealESRGAN degradation pipeline has become widely adopted in recent years, and multiple models have shown its satisfying generalization to real-world images. Perpendicular efforts have tackled issues like architectural improvements or the use of generative priors. The proposed method above focuses on blind SR.

[110] Diffusion for super-resolution: The impressive performance of SD for generation tasks resulted in early interest in the SR community. Some techniques adopted SD for SR in some constrained settings like face and class-constrained SR, while a more general formulation for in-the-wild images was later proposed. Other approaches leverage the image priors of a pretrained text-to-image SD model for SR. While leveraging GAN image priors is well-established in the literature, their lack of generality restricted their use to favorable settings, e.g. face super-resolution. Instead, SD has shown remarkable general-purpose image generation capabilities. Thus, several recent works have exploited SD priors for SR. Some use adaptors on top of a frozen pre-trained text-to-image model, achieving remarkable visual quality for blind super-resolution. However, simple concatenation has been used which in contrast to the bidirectional encoding above does not allow the training of an efficient architecture. Our proposed method also benefits from scale distillation to further improve visual quality.

[111] Efficiency for SD: Improving efficiency has been a major focus within the SD literature. There are two predominant approaches for making diffusion models more efficient. A line of work focuses on reducing the number of inference steps including more effective sampling strategies and distillation techniques. A second line of work focuses on making the SD model more lightweight but often still uses multiple steps for inference. In proposed method Edge-SD-SR, a lightweight encoder and decoder is used. Overall, there is no prior art work which both reduces inference steps and deploys a lightweight model for the whole architecture (U-Net, encoder / decoder). To enable this, the solution proposed above is tailored to the task of SR, by introducing the newly proposed bidirectional conditioning.

[112] The proposed method uses a general formulation where the noise variance is also conditioned on the LR image. As shown above, conditioning noise variance plays a vital role in performance. The proposed method has a training paradigm for the encoder training and shows that the encoder size can be effectively reduced for an on-device model. The proposed method is an SD based SR solution that can readily run on device. The complete solution proposed above tackles speed and size to make the diffusion model more efficient. To enable training, a new conditioning technique which provides a better training signal for the diffusion model to more directly transition from LR to HR. The proposed method can be combined with other techniques specifically targeting improving inference speed of SD-based SR to yield an efficient and fast model.

[113] Those skilled in the art will appreciate that while the foregoing has described what is considered to be the best mode and where appropriate other modes of performing present techniques, the present techniques should not be limited to the specific configurations and methods disclosed in this description of the preferred embodiment. Those skilled in the art will 5 recognise that present techniques have a broad range of applications, and that the embodiments may take a wide range of modifications without departing from any inventive concept as defined in the appended claims.

Claims

1. A computer-implemented method for training, on a server, a diffusion machine learning, ML, model for increasing a resolution of an input low-resolution, LR, image, wherein the diffusion ML model comprises a LR encoder for generating a LR embedding vector for the input image, a denoising ML model and a decoder, the method comprising:obtaining a training dataset comprising at least one high resolution, HR, image;inputting a HR image from the training dataset into a HR encoder to generate a HR embedding vector for the HR image, wherein the HR encoder is separate from the diffusion ML model which is being trained;obtaining a LR image from the HR image;inputting the obtained LR image into the LR encoder to generate the LR embedding vector for the input LR image, wherein the LR embedding vector comprises a mean embedding vector and a variance embedding vector for the input LR image;obtaining, at sequential time steps, a plurality of sample latent variables by combining the HR embedding vector with a noise sample taken from a Gaussian distribution defined by the generated mean and variance embedding vectors,predicting, for each sample latent variable, using the denoising ML model, a HR embedding vector; andtraining the denoising ML model by minimising a difference between the predicted HR embedding vectors and the generated HR embedding vector; andoutputting, from the server to a user device, the trained diffusion ML model which comprises the LR encoder, the trained denoising model, and the decoder.

2. The method of claim 1, wherein training the denoising ML model comprises checking whether parameters of the denoising ML model are converging at each sequential time step and when there is no convergence, inputting a further HR image from the training dataset and repeating the steps of obtaining a LR image, inputting the obtained LR image, obtaining a plurality of samples, predicting and training until there is convergence.

3. The method of claim 1 or claim 2, wherein combining the HR embedding vector with the noise sample comprises using a weighted sum of the noise sample and the HR embedding vector.

4. The method of claim 3, wherein using the weighted sum comprises using a first weight for the noise sample and a second weight from the HR embedding vector, wherein the second weight increases relative to the first weight with each sequential time step during a denoising process of the denoising ML model.

5. The method of any one of the preceding claims, further comprisingtraining the LR encoder whilst keeping parameters of the HR encoder frozen, and outputting, from the server to the user device, the trained diffusion model comprising the trained LR encoder and the trained denoising model.

6. The method of claim 5, further comprising using a scheduler to schedule training of the LR encoder jointly with the denoising model for at least some of the sequential time steps, wherein jointly training the LR encoder and the denoising model comprises adjusting parameters of both the LR encoder and the denoising model to minimise a difference between the predicted HR embedding vectors and the generated HR embedding vector.

7. The method of claim 6, further comprising using the scheduler to schedule training of only the denoising model for other time steps, wherein the parameters of the LR encoder are frozen and the method comprises adjusting only parameters of the denoising model at the other time steps to minimise a difference between the predicted HR embedding vectors and the generated HR embedding vector.

8. The method of claim 6 or claim 7, further comprising using the scheduler to decrease a time gap between successive sequential time steps when scheduling training of the LR encoder jointly with the denoising model.

9. The method of claim 8, further comprising using the scheduler to increase a time gap between successive sequential time steps until the time gap converges to a uniform time gap for later time steps when scheduling training only of the denoising model.

10. The method of any one of the preceding claims, further comprising: training the decoder before outputting the trained diffusion model and outputting the trained diffusion model with the trained decoder.

11. The method of any one of the preceding claims, further comprising using scale distillation.

12. The method of any one of the preceding claims, further comprising reducing a size of the denoising model before training the denoising model.

13. The method of claim 12, further comprising reducing the denoising model to a single layer.

14. The method of any one of the preceding claims, further comprising reducing a size of the LR encoder before training the denoising model.

15. A computer-implemented method for generating, on a user device, a high resolution, HR, image from a low resolution, LR, image, the method comprising:receiving an LR image;inputting the received LR image into a trained diffusion ML model on the user device, wherein the diffusion ML model has been trained according to anyone of the preceding claims and comprises an LR encoder, a trained denoising model and a decoder;generating, using the LR encoder, a LR embedding of the received LR image;inputting the LR embedding into the trained denoising ML model to generate a HR embedding; andgenerating, using the decoder, a HR image from the generated HR embedding.

16. A computer-readable storage medium comprising instructions which, when executed by a processor, causes the processor to carry out the method of claims 1 to 15.

17. A server for training a diffusion machine learning, ML, model for increasing a resolution of an input low-resolution, LR, image, the server comprising:at least one processor which is configured, individually and / or collectively, to carry out the method of any one of claims 1 to 14.

18. A user device for generating a high resolution, HR, image from a low resolution, LR, image, the user device comprisingat least one processor which is configured, individually and / or collectively, to carry out the method of claim 15.

Images