With the advent of machine learning, most of the state of the art solutions in image processing tasks have become datadriven: Many degrees of freedom allow a trainable model to adapt well to the input data. While yielding highly tailored solutions, results are often too complex for a rigorous analysis. In contrast to this, classical model-driven approaches provide a sound theoretical foundation but often cannot compete with their learning-based counterparts.
The goal of this work is to combine the advantages of model-and data-driven approaches on the basis of wavelet shrinkage for denoising: Given a signal that has been corrupted by noise, the goal is to reconstruct the original image as well as possible by modifying parts of the signal in the wavelet domain. We exploit the flexibility of a learning-based approach to train a smooth shrinkage function that adapts both to the wavelet scales and to the noise level. The proposed shrinkage function can even amplify wavelet coefficients and thus enhance important image structures, a property that most This work has received funding from the European Research Council (ERC) under the European Union's Horizon 2020 research and innovation programme (grant agreement no. 741215, ERC Advanced Grant INCOVID). established shrinkage functions do not share. A low number of trainable parameters allows us to manually inspect the learned results and infer smooth connections between them. From these, we design a generic compact shrinkage function that incorporates the learned adaptivity, while only using two parameters. Experiments show that our shrinkage function outperforms classical ones by a significant margin. To the best of our knowledge, this is the first work to present a learningbased shrinkage function with this level of compactness.
Wavelet shrinkage has first been proposed by Donoho and Johnstone [1] . Therein, a noisy signal is transformed to a wavelet basis, resulting coefficients are modified by a shrinkage function and then transformed back to obtain a denoised result. Classical shrinkage functions use a single threshold parameter for all scales of the wavelet transformation. However, individual scales might require distinct thresholds as they are differently affected by noise. To this end, adaptive shrinkage methods have been proposed. Early works of Zhang et al. [2, 3] already study a smooth adaptive shrinkage function with trainable thresholds. Other authors train arbitrarily shaped shrinkage functions [4] , some also train the wavelets [5] or even general adaptive filters for shrinkage operations [6] . Non-trainable adaptive statistical models include [7, 8] . Most learning-based methods produce large amounts of trained parameters while relations between them are rarely investigated. In contrast to this, we directly employ a tightly parametrised model from which it is easy to infer smooth underlying parameter relations.
An important connection between 2D wavelet shrinkage and nonlinear diffusion filtering has been established by Mrázek and Weickert [9] . It allows us to directly translate a diffusivity into a trainable shrinkage function. We use a so-called Forward-and-Backward (FAB) diffusivity, resulting in a shrinkage function that can amplify coefficients. Only few works such employ this property directly [10] , but also results for learned shrinkage functions suggest its usefulness [4] . The corresponding concept of backward diffusion has also shown to be successful [11, 12, 13] .
The remainder of this paper is structured as follows: In Section 2 we review classical wavelet shrinkage. We propose our model in Section 3, which is experimentally evaluated in Section 4. Finally, Section 5 summarises our conclusions. 
Classical discrete wavelet shrinkage represents a noisy signal f = v + n in the wavelet basis, wherein the additive noise n is better separated from the true signal v. This is achieved by the following three-step framework:
1. Analysis: The input data f is transformed to wavelet and scaling coefficients. In this representation, the noise affects all wavelet coefficients while the signal is represented by only a few significant ones [14] .
2. Shrinkage: A shrinkage function S θ with a threshold parameter θ is applied individually to the wavelet coefficients. The scaling coefficients remain unchanged.
The denoised version u of f is obtained by back-transforming the modified wavelet coefficients.
Many shrinkage functions have been proposed. We will consider the most prominent ones of soft [15] , hard [14] and garrote [16] shrinkage for comparison. While these functions are easy to use in a practical setting, they suffer from applying the same threshold parameter to all scales and their binary decision structure. Finer scales might require a different threshold than coarser scales as they are more affected by noise. Furthermore, there is no clear separation between noise and signal coefficients such that eliminating too many coefficients always destroys signal details and eliminating too few leaves noise in the reconstruction. The classical wavelet transformation is not shift-invariant: Shifting the input signal f will produce a different set of wavelet coefficients. To overcome this problem, Coifman and Donoho suggested cycle spinning [17] : The input signal is shifted, wavelet shrinkage is applied and the results are averaged for all possible shifts. This yields the shift-invariant but redundant non-decimating wavelet transformation.
For the two-dimensional wavelet transform, wavelet coefficient channels w x , w y and w xy for x-, yand diagonal direction are obtained, and special care is required to design rotationally invariant shrinkage rules. Mrázek and Weickert [9] achieve this by a channel coupling that is inspired by nonlinear diffusion filtering. Their Haar wavelet shrinkage rule is given by
The argument w 2 x + w 2 y + 2 w 2 xy is a consistent approximation to the rotationally invariant expression |∇u| 2 , where | . | is the L 2 norm and ∇ = (∂ x , ∂ y ) denotes the spatial gradient operator. The function g is a diffusivity function from a nonlinear diffusion filter [18] . In nonlinear diffusion, filtered versions u(x, t) of an image f (x) are obtained by solving the partial differential equation
with u(x, 0) = f (x) as initial condition and diffusion time t.
The diffusivity g steers the activity of the process. It is scalarvalued and becomes small at edges where |∇u| 2 is large. This results in rotationally invariant edge-preserving denoising.
Mrázek and Weickert [9] have shown that one explicit time step of nonlinear diffusion with diffusivity g is equivalent to the Haar wavelet shrinkage (1). Thus, we can directly translate existing diffusivities to shrinkage functions.
In our model, we equip the two-dimensional coupled Haar wavelet shrinkage approach from [9] with a trainable adaptive shrinkage function. We sample a noisy image f on a discrete grid with 2 L × 2 L sampling positions and grid distance h x = h y =: h. The discrete image is represented by a vector f . It is transformed with the non-decimating two-dimensional Haar wavelet transformation W into lowpass scaling coefficients at scale L denoted by w L , and directional wavelet coefficients w x , w y , w xy L =1 for different scales. To ensure that coefficients have the same range on all scales, we rescale the basis functions accordingly. A shrinkage function S θ with parameters θ = θ 1 , . . . , θ L for each scale is applied to the wavelet coefficients. Coefficients on scale are modified component-wise by S θ according to the coupled shrinkage rule (1) . We obtain the reconstruction u by applying the backward transformationW to the modified set of wavelet coefficients and the unaltered scaling coefficients:
We found the Forward-and-Backward (FAB) diffusivity of Smolka [19] to be a good candidate for modelling our shrinkage function. It uses two contrast parameters λ 1 and λ 2 that control the amount of forward and backward diffusion. The diffusivity is given by
It is translated into the shrinkage function according to the coupled shrinkage rule (1) . For the extreme case of λ 1 = λ 2 , the diffusivity simplifies to the exponential Perona-Malik diffusivity [18] , corresponding to pure shrinkage. For larger differences between λ 2 and λ 1 , the backward diffusion becomes more pronounced and the shrinkage function damps small coefficients and amplifies larger ones. 
To train the shrinkage functions, we add Gaussian noise to a database of ground truth images (v k ) K k=1 . This yields noisy images f k from which we compute denoised results u k . The trainable parameters for the proposed shrinkage function on scale are given by θ = λ 1 , λ 2 . As an objective function, we choose the mean square error between the reconstruction u k and the corresponding ground truth image v k , averaged over all image pairs and normalised by the number of pixels:
The parameters are trained by the gradient-based L-BFGS algorithm [20] . To that end, we compute the gradients of the objective function w.r.t. all trainable parameters.
In our experimental setup, we use 400 images from the BSDS500 database [21] as a training set and the 68 images introduced in [22] as a test set. Their grey values are in [0, 255] and the grid size is set to h = 1. From each image we select random regions of size 256 × 256, i.e. the number of scales is L = 8. All images are corrupted by Gaussian noise of mean 0 and standard deviation σ. We do not clamp the resulting pixel grey values to the original grey value range to preserve the Gaussian statistics of the noise. We have found the learned parameters to be robust w.r.t. any reasonable initialisation, so no pretraining is performed.
In a first experiment, we train the adaptive shrinkage function for σ = 25. The results of the learned shrinkage functions are presented in Figure 1 . On the first and finest wavelet scale, all coefficients are shrunken. On the second scale, we observe coefficient amplification. We presume that this compensates the loss of image details caused by shrinking coefficients on the first scale. All further scales do not perform significant shrinkage as the learned function approaches the identity. Therefore, we do not display coarser scales > 4. When we increase the noise level to σ = 50, we observe that more scales are involved in the shrinkage process. The trained results are displayed in Figure 2 . Both configurations are in line with our conjecture that the diffusivity should change smoothly over scales and noise level. With coarser scales, the amount of shrinkage and amplification decreases, i.e. λ 1 and λ 2 tend to 0. With increasing noise, shrinkage and amplification become stronger and cover more scales.
To investigate the effectiveness of different aspects of our model, we perform the following ablation study. We start with classical hard wavelet shrinkage and equip it with the nondecimating wavelet transformation and the coupled shrinkage rule to enable a fair comparison. For σ = 25, we obtain a test MSE of 156.8. In a second step, we use the proposed shrinkage function restricted to λ 2 = λ 1 , so no amplification can take place. Still, the shrinkage function does not adapt to the individual scales. This marginally improves the test MSE to 156.3, showing that smoothness of the shrinkage function alone does not matter for reconstruction quality. When we remove the restriction on the shrinkage function, the MSE drops down to 148.7, which indicates that the amplification of wavelet coefficients is helpful for a good reconstruction. Finally, introducing adaptivity to the scales boosts the MSE further to 132.2, demonstrating that the scale dynamic is the crucial ingredient for a good denoising result.
So far, the contrast parameters are trained for each pair of scale and noise level σ from which we will now infer a generic relation. Figure 3 shows the evolution of both contrast parameters over the scales and the noise levels. A function of type α 2 with an appropriate scalar α can provide a good description of the scale dependence of λ 1 . For λ 2 , the values on fine scales do not follow this relation, but in these cases all relevant coefficients are already covered by shrinkage through a large λ 1 , making the choice of λ 2 irrelevant.
Regarding the relationship between the shrinkage functions and the noise standard deviation it was already noted in [4] that a simple rescaling of shrinkage functions is sufficient for adapting to a new noise level. For our parametrisation, rescaling the complete shrinkage function is equivalent to rescaling both λ 1 and λ 2 . In Figure 3 we can see that indeed such a rescaling is learned. These two insights suggest that a suitable generalisation of the shrinkage function parameters which is smooth over the scales and the noise standard deviation σ is given by λ 1 ( , σ) = ασ 2 and λ 2 ( , σ) = βσ 2 where α and β are scalars that have yet to be determined. To empirically show that this parametrisation indeed captures the underlying relations in a reasonable way, we compare two models: One model trains the proposed shrinkage function for each pair ( , σ) individually, while the other one only optimises the factors α and β of the generalised parameters. To ensure a fair comparison, both models are trained on a new training and test set combining images with noise levels between σ = 10 and σ = 60 in steps of 2.5. Indeed, the generic model performs only 5.1% worse than the model with individual parameters in terms of MSE, while training only 2 instead of 336 parameters. With this result we conclude that the generic shrinkage function captures the adaptivity to scales and noise levels well. The scalars are learned as α = 5.4 and β = 8.9, yielding a combined generic coupled shrinkage function (1) 
Lastly, we compare our generic shrinkage function to soft, hard and garrote shrinkage over a range of noise levels. We optimise the threshold parameter of the classical functions individually for each noise level, while the generic function is used as is from (6). The results are displayed in Figure 4 . Although the classical approaches are optimised for each noise level, they are inferior to the generic shrinkage function. For higher noise levels, the margin becomes larger as the generic function benefits more from image enhancement through coefficient amplification. For σ = 50, Figure 5 shows reconstructions along with the noisy input and the ground truth image. Soft shrinkage blurs the image too strongly since all wavelet coefficients are decreased by the same margin. Hard shrinkage suffers from remaining noise as it does not shrink large noisy wavelet coefficients. While less pronounced, this is also the case for garrote shrinkage. Both garrote and hard shrinkage also blur important image structures. Our generic shrinkage function outperforms all classical approaches. By strongly shrinking coefficients on fine scales, noise is efficiently removed. To compensate for lost image details, amplification of wavelet coefficients on coarser scales enhances important image details.
Our approach of learning a compact shrinkage function for wavelet denoising combines the advantages of model-driven and data-driven approaches: In contrast to other parameter learning strategies, we can cope with as little as two parameters without substantially sacrificing performance. This results in an interpretable shrinkage function and a transparent, but adaptive model.
In our ongoing work we extend these findings to other adaptive nonlinear approaches such as diffusion evolutions.