Pivotal Auto-Encoder via Self-Normalizing ReLU

We begin by introducing a framework parallel to the analysis sparse representation model, presented in Section I, named the synthesis model. This model serves as the premise for introducing the square root lasso algorithm, which we will extend to the transform model. The synthesis model assumes that a signal $x\in\mathbb{R}^{n}$ can be represented as a linear combination of a few columns, called atoms, from a matrix $D\in\mathbb{R}^{n\times d}$ named dictionary. In plain language, the signal corresponds to a multiplication of a dictionary by a sparse vector $z^{*}\in\mathbb{R}^{d}$, i.e.,

Various algorithms have been proposed to implement the sparse coding task of estimating $z^{*}$ given $x$ [17]. These include the matching pursuit [18] and the basis pursuit [19], also called lasso in the statistics literature [20]

As with (2), the primary drawback of lasso is that the optimal value of the parameter $\lambda_{L}$ is dependent on the noise level and, therefore, must be adjusted for each specific noise level. For example, in a Gaussian noise environment, $\lambda_{L}$ proportional to $\sigma\sqrt{\log d/n}$ is minimax optimal for signal reconstruction, $\|\widehat{x}-x\|_{2}$, in high dimensions [21]. Therefore, to achieve a good estimation of $z^{*}$ or prediction of $x$ for an unknown noise level, one must estimate $\sigma$.

Square root lasso. A modified version of the lasso has been proposed to solve the dependence of $\lambda$ on noise power, called square root lasso [15]

which takes the square root of the error term of (4). Belloni et al. [15] have proven that the square root lasso achieves minimax optimality of estimation and signal reconstruction error for a hyperparameter $\lambda_{\sqrt{L}}$, which is pivotal to the noise level. For instance, in the case of Gaussian noise, the minimax optimal $\lambda_{\sqrt{L}}$ is proportional to $\sqrt{\log d/n}$, which is independent of $\sigma$ and leads to a constant parameter for all Gaussian distributions. Moreover, numerous algorithms have been developed to efficiently solve the problem based on its convex property [21].

Although the square root lasso is powerful and attractive, it has never been applied in the context of dictionary learning or adapted to form a novel neural network architecture. In this work, we draw an exciting connection between the square root lasso and transform modeling. The extension to synthesis model and dictionary learning is a direct consequence of the analysis of the transform learning since (3) can be seen as a particular case of (5), where $D=I$ and the input signal is $Wy$. Further details are provided in Appendix B.

As elaborated in Section I, the transform model and sparse auto-encoders are tightly connected. The forward pass in sparse encoders essentially acts as a sparse representation pursuit within the transform model. In this framework, the transformation is characterized by the weights within a set of layers, which may include convolutional ones. Thus, the transformation is also learned during the model training process. This method of deriving the transformation directly from the data is known as transform learning.

At its core, transform learning [22] employs a data-driven feature extractor to transform input data into a suitable representation. This approach can improve upon the limited ability of analytical transformation methods, such as wavelets, to handle data. As a result, transform learning often yields superior restoration performance compared to the analytical approaches [6].

Transform learning is comparable to dictionary learning from an analysis perspective [23]. In dictionary learning, a basis $D$ is trained to recover the data, $x=Dz^{*}$, from the representation $z^{*}$ [24]. In contrast, transform learning aims to learn a transformation $W$ to generate the representation $z^{*}$. Since the transformation $W$ is data-dependent and the output representation $z^{*}$ is unknown, jointly learning both is a challenging task.

In this work, we propose a new learning problem that combines the transform model with deep learning through the interpretable framework of transform learning. Our objective is to demonstrate that our method can exhibit the interesting properties of the classical problem established in Section III. We extend our algorithm to the transform learning framework and demonstrate its effectiveness in enhancing the robustness of deep learning through experiments in Section V.

Blind denoising is the task of removing noise from an input signal when the noise magnitude is unknown at test time. This task is closely related to our objective. Several methods have been proposed to tackle this task, with the most common approach for blind denoising being to estimate the noise distribution (or simply the noise level) to identify and remove it from the signal [25, 26, 27]. However, this method lacks flexibility because it requires learning different weights for each noise level and solving the difficult problem of noise estimation.

In contrast to this approach, where all weights are learned from scratch for each noise level, some existing methods recognize that not all weights depend on the noise and only adjust the regularization parameters, such as the bias in the case of auto-encoders, for different noise levels [14]. This method reduces the overall number of network parameters to be learned, but they still need to be adjusted for each noise level.

Another standard technique, which emerged alongside the advancement of deep learning in computer vision tasks, is to train one model across a wide range of expected noise levels [28]. In this case, the denoising performance of such a model is generally inferior compared to a model trained for a specific noise distribution [29]. Moreover, it has been shown that a model trained using this method tends to focus on the average noise level of the training range, rather than learning generalizable weights for all noise levels. To address this issue, Gnansambandam et al. [29] proposed determining the optimal noise training sample distribution from a minimax risk optimization perspective. The approach proposed in [29] is orthogonal to ours, as it does not suggest modifying the network architecture but instead focuses solely on the training strategy.

In this work, we propose a new architecture for implementing (3) that is inherently noise level independent. Our theoretical study, presented in Section III, shows that the same model parameters achieve high-quality signal recovery across all noise levels when learned correctly. Indeed, our experimental results indicate that a single neural network can be trained and practically applied to handle all noise levels, without re-training or updating the bias term.

First, we prove the recovery of the correct support in the presence of bounded noise, a prevalent scenario in the robustness literature [30, 31]. Subsequently, we extend these results to Gaussian noise within a probabilistic setting. To prove these results, the following assumptions are necessary.

On a different note, we can deduce that the correct support can be fully recovered if the signal-to-noise ratio is sufficiently high, as formally stated in Theorem 1. Furthermore, $\lambda$ can be conveniently bounded in all cases by

A bound for $\lambda$ can be used to guarantee that no false positive support error occurs. Improved bounds can be achieved with additional assumptions on the noise. In fact, we investigate this aspect for the Gaussian case in Section III-C.

We now proceed to show that the estimation error, $\|\widehat{z}-z^{*}\|_{2}$, can also be bounded with the same choice of parameter $\lambda$, under the same Assumptions 1 and 2.

In Appendix F, we present an empirical validation of Theorem 2.

We now extend the results of Section III-A to the Gaussian case. The key point of this analysis is that we can use practical values for $\lambda$, which can be computed independently of the noise level. We show that these $\lambda$ values are suitable for any additive Gaussian noise in the input and are thus pivotal to its standard deviation $\sigma$.

In this section, we introduce an iterative optimization algorithm for minimizing (3) that can be efficiently implemented and formulated as a novel sparse auto-encoder architecture. It is worth noting that this objective function corresponds to a convex optimization problem. Therefore, it inherits not only all the theoretical properties of convex optimization problems, but also the algorithms that can be used to solve it, such as the interior point method [34] or the alternating direction method of multipliers [35].

In [21], the authors studied the geometric structure of the square root lasso problem and concluded that the $\ell_{2}$ loss function is non-differentiable only in extreme cases of overfitting. In practice, this situation is rare when the data are corrupted by noise and a sufficiently large regularization parameter $\lambda$ is used to produce a sparse solution. Consequently, the data fitting term in the objective function behaves as a strongly convex and smooth function.

Leveraging these attractive geometric properties, we can use proximal gradient descent to iteratively minimize (3). The theoretical analysis in [21] shows that such an optimization algorithm achieves fast local linear convergence. The same theoretical justification also applies to (3) since our optimization problem can be viewed as a simpler instance of the square root lasso, presented in Section II-A. Therefore, we propose adapting proximal gradient descent to the problem studied here, as described in Algorithm 1, which we have named Self Normalizing ReLU or NeLU for short.

The algorithm works by continuously refining the solution via iterative updates, progressing in the direction opposite to the gradient of the objective function. Each iterative update incorporates a proximal operator, which introduces a penalty term to the objective function, thereby promoting sparsity. In the specific problem at hand, the proximal operator takes the form of a soft-thresholding operator, as outlined in Section I. This operator performs a shrinkage operation on the variables, setting any values below a specified absolute threshold to zero. Importantly, the soft-thresholding operator can be replaced with ReLU to enforce nonnegative representations. The iterative process continues until the desired level of convergence is achieved.

We propose to adapt Algorithm 1 for transform learning by unrolling the algorithm into a layered neural network architecture, following the approach presented in [36]. The idea is to unfold an iterative algorithm and construct it as a network architecture, mapping each iteration to a single operation and stacking a finite number of operations on top of each other. This approach enables the incorporation of a wide range of mathematical techniques into deep learning models [37, 38, 39]. Specifically, we achieve this unfolding by limiting the number of iterations in Algorithm 1 to $N$ iterations. The resulting architecture is depicted in Figure 1.

In addition, we propose to improve the performance of the algorithm by employing Nesterov acceleration [40]. Nesterov acceleration is a variant of momentum that speeds up the convergence of gradient descent algorithms, and has demonstrated efficacy in various contexts. By incorporating it into Algorithm 1, we aim to achieve superior performance in transform learning tasks, based on the understanding that accelerated gradient descent converges faster and operates effectively with shallower networks, which are easier to train.

Finally, we note that all parameters of the resulting algorithm, including $W$, $\lambda$, $\beta$, and $\alpha$, can be trained end-to-end. This means that the network can be trained on a dataset to learn the optimal values for these parameters, allowing it to perform well on various transform learning tasks. The final accelerated algorithm is presented in Algorithm 2.

At each iteration, the algorithm computes the gradient of the objective function with respect to the model parameters at a point in the direction of the momentum and updates the momentum in the opposite direction of the gradient. The solution is then updated based to the momentum, taking into account the proximal operator. This operator introduces regularization to prevent overfitting and improve the generalization performance of the model.

First, we present experiments conducted on synthetic data to demonstrate the advantage of our proposed method over the traditional sparse encoder algorithm (2). We assume that the transformation $W$ is known and construct a $100\times 100$ random matrix, where each entry is sampled from the standard normal distribution, and then normalize the rows to have unit $\ell_{2}$ norm to satisfy Assumption 4. Next, we generate the input signal by creating a vector $z^{*}$ with fixed sparsity level, following the procedure described in [41], to obtain a signal consistent with the transform model, such that $\|Wx\|_{0}=5$. Finally, we contaminate the signal with i.i.d. Gaussian noise of level $\sigma$ to produce the measurements $y=x+\xi$.

We evaluate the estimation error, $\|\widehat{z}-z^{*}\|_{2}$, which measures the distance between the estimated signal $\widehat{z}$ and the true signal $z^{*}$, as a function of the noise standard deviation $\sigma$. We compare the solutions obtained by minimizing (2) and (3) in two settings. In the first setting, we perform an oracle cross-validation that sweeps over a range of parameters $\lambda$ to find the regularization parameter that minimizes the estimation error for each algorithm. Importantly, this setting is infeasible since it requires the ground truth data to calculate the estimation error. Nevertheless, it reveals the best performance that one can hope to achieve. In the second setting, we consider a more realistic scenario where we compare the performance of the proposed Algorithm 1 using the theoretical value of $\lambda$ divided by 2, following Belloni’s empirical improvement [15].

Figure 2 shows that both algorithms achieve similar performance in the oracle setting. However, the optimal value of $\lambda$ for the soft-thresholding algorithm (2) is linearly proportional to the noise level $\sigma$, while the optimal value for the proposed algorithm (3) is pivotal to it. This conclusion is consistent with our theoretical analysis presented in Section III. Additionally, we observe that the theoretical value of $\lambda$ for the proposed algorithm, described in Theorem 1, is very close to the actual optimal value. This indicates that the proposed optimization problem may be a reasonable alternative to the current soft-thresholding algorithm, which forms classic encoder architectures, in practical situations where the noise level is unknown.

Trainable transforms often outperform analytical transformations, such as total variation and wavelets, in most signal processing applications [6]. This is because the transformation used in signal processing is frequently unknown and must be inferred from the data. This motivates the use of neural networks to simultaneously learn the transformation and the sparse representation of the data.

The first learning task is supervised sparse coding, where the input consists of signals $y$ determined by the model and their corresponding synthetic sparse vectors $z^{*}$ generated by a sparsifying transform, as described in Section V-A. In this case, the goal of the neural network is to learn the transformation $W$ stored in its weights and accurately identify the corresponding sparse output vector given a set of input-output pairs. Mathematically, this can be expressed as an end-to-end training scheme, minimizing the cost function

where $\widehat{z}$ is the output of the network outlined in Figure 1.

To compare the effectiveness of the proposed approach, two different neural network architectures are used: one based on Algorithm 2, and a baseline sparse encoder architecture. Both networks consist of a linear layer followed by a thresholding layer. In the baseline version, the thresholding layer is the soft-thresholding non-linearity. In contrast, our proposed architecture uses the Accelerated NeLU non-linearity as presented in Algorithm 2, with the soft-thresholding operator. Both networks are trained using the AdamW optimizer [42] and minimize the mean squared error (MSE) loss with a fixed noise standard deviation of $\sigma=0.1$. After training, we evaluate the performance of the fitted networks on different noise levels, $\sigma$, than those used during training. The results, displayed in Figure 3, suggest that NeLU is significantly more robust to unseen noise levels than soft-thresholding, indicating that the NeLU non-linearity results in a predictive model that generalizes to other noise levels without additional training.

We also demonstrate the effectiveness of the proposed model in signal denoising applications. In this task, the input is a noisy signal $y=x+\xi$ and the goal is to produce a cleaned version of the signal, $\widehat{x}=W^{+}\widehat{z}$, by removing the additive noise. To this end, a final linear layer is added to each network according to the sparse model. The first two layers perform sparse coding, i.e., they estimate the sparse representation, while the last layer projects the sparse estimation back into the input space. Sparse data is generated in the same manner as before and the MSE loss is minimized. The results, shown in Figure 4, again reveal that the NeLU leads to more stable recovery than the soft-thresholding non-linearity.

In this experiment, the aforementioned networks are employed to perform natural image denoising using a patch averaging technique based on the Convolution Sparse Coding model [43]. To accomplish this, each input image is replicated stride times and translated across every dimension, producing slight shifts of the original for each replica. As a result, a single image transforms into a set of $\textit{stride}^{2}$ slightly offset variations. These shifted versions of the input are then processed collectively by the network, resulting in intermediate (shifted) denoised versions of the same input image. Finally, these intermediate denoised output images are shifted back and averaged to yield the final reconstructed output image. This process is visualized in Figure 5 for a one-dimensional (1D) signal. For more details, see [43].

The datasets and preprocessing procedures are adopted from [43]. Specifically, we use clean training images from the Waterloo Exploration dataset [44], and a validation set consisting of $432$ images from BSD [45]. Noisy images are generated by adding white Gaussian noise with a constant standard deviation $\sigma=15$. In each iteration, we randomly crop a patch of size $128^{2}$ from an image and obtain a random realization of the noise.

In this setting, we replace the linear layers with convolution and deconvolution layers, respectively. Concurrently, the soft-thresholding operator is substituted by ReLU as the proximal operator. The models learn 175 filters of dimensions $11\times 11$ and a stride of 8. We utilize the AdamW optimizer with a learning rate of $2\cdot 10^{-2}$, which is reduced by a factor of 0.7 after every 50 epochs. Additionally, the optimizer’s $\epsilon$ parameter is set to $10^{-3}$ to ensure stability. The models are trained for 300 epochs, minimizing the MSE loss.

To evaluate the performance of the models, we use the BSD68 dataset, which is distinct from the validation set. The experimental results, as shown in Table I, allow us to compare the performance of each model on the test dataset at varying noise levels. We can see that the proposed Algorithm 2 layer outperforms the ReLU activation function for virtually all noise levels, and the performance gap widens as the noise level deviates further from the trained noise level.