Structural Restricted Boltzmann Machine for image denoising and classification

RBMs were originally proposed for generative purposes, where the probability distribution of a given dataset is approximated, from which samples could be drawn. When training each model, the mean LL of the train and validation set are evaluated. This is indeed the objective function the model optimises using mini-batch SGD with CD. For a set of visible data $\mathcal{D}$, the LL is computed as follows

where the Annealed Importance Sampling (AIS) method is used to approximate the $\log Z(\theta)$ term. We follow the setting presented in [14] to estimate the partition function of the RBMs with 1000 AIS iterations, although using 5 times fewer intermediate distributions.

One of the problems these models can handle is image denoising, consequently, we assess the performance of our models on this task as well. Given an image that is the distorted version of an original image, the denoising problem consists of recovering the original image by removing noise. RBMs and SBMs trained to approximate the underlying distribution of a set of images can be used for this purpose, by means of Gibbs sampling. Images are reconstructed after a single Gibbs sampling step by setting the corrupted image as initial state, the resulting visible state is taken as the reconstructed image. The performance is measured by computing the mean square error (MSE) between the reconstructed image and the original non-corrupted image.

Even though RBMs were originally introduced as generative models, they can be used for classification. RBMs are commonly trained to extract features, which are then fed to a classifier network [1], this can be seen as a feature extraction method. Nonetheless, RBMs can also be used as stand-alone discriminative models, as proposed by Larochelle and Bengio in [15]. In [16], authors use a convolutional deep BM to classify Electroencephalograms for early Alzheimer’s disease diagnosis, and in [17] a two-layer deep BM which classifies facial expressions from thermal infrared images is proposed.

In this work, the Classification RBM proposed by Larochelee and Bengio in [15] is used to perform image classification purposes. In this model, classes are represented with a one-hot encoded vector $\textbf{y}=(1_{y=k})_{k=1}^{C}$, and they are added to the RBM as additional visible units. These variables are connected to the hidden units via the matrix $U$ and have biases c. Consequently, similarly to Eq. (1), the energy of this model is

Which leads to the modification of the expressions of some conditional probabilities for GS

Deriving how to update the new parameters using CD is straightforward. However, an alternative method more suited for classification was suggested in [15].

The probability of each class $y_{k}$ given an instance v can be computed via

Authors in [15] proposed to train the model such that the probability of the true class given an image is maximised. Hence, the discriminative objective function can be written as in Eq. (16).

The gradients of this objective function with respect to the parameters can be computed exactly as given in Eq. (17)

where $o_{kj}(\textbf{v})=b_{j}+U_{kj}+\sum_{i}W_{ij}v_{i}$. Therefore, when training our models for the classification problem, Eq. (17) is used when updating the parameters with SGD.

Standard metrics to evaluate a classifier are used to assess classification RBMs as well. Classification accuracy is one of the most popular metrics, which evaluates the ratio between the number of well-classified instances and the total number of instances classified. This metric has the risk of not being representative of the classification performance when classes are unbalanced. Therefore, when classes are unbalanced, the balanced version of accuracy can be used, where the contribution of each class is weighted with the inverse of its frequency of appearance.

Nonetheless, since our models are able to return the probability of each class given an image, other metrics that penalise giving lower probabilities to the correct class can be used. One of these metrics is the Log-loss function

where $Y$ is a vector of $N$ one-hot represented $K$ classes, so $y_{i,k}=1$ if the instance $i$ corresponds to class $k$ and 0 otherwise, and $P$ is the vector of probabilities the model assigns to each instance. The closer the Log-loss value is to 0, the more accurate and confident the predictions of the model are. In this case, the same strategy used with the accuracy is used when classes are unbalanced, assigning weights to each class according to the inverse of its frequency of appearance.

The introduced SBM models take advantage of the underlying structure of the problem to reduce the number of trainable parameters. As an area of application, we focus on images in this work. We assume that in the image problems addressed there is a spatial relationship between the value of a pixel and its neighbourhood. Hence, each hidden unit may focus on different patches of visible units (pixels), and thus, each patch/neighbourhood of visible units will take the role of $\mathcal{V}(h_{j})$. In order to define these sets, given a pixel $\textbf{u}=(u_{x},u_{y})\in\mathcal{I}$ within a $d\times d$ image, we use the Chebyshev distance between two pixels $d(\textbf{u},\textbf{u}^{\prime})=\max(|u_{x}-u^{\prime}_{x}|,|u_{y}-u^{\prime}_{y}|)$ to define the neighbourhood. The neighbourhood of a pixel $\textbf{u}^{\prime}$ is defined as $\mathcal{N}(\textbf{u}^{\prime})=\{\textbf{u}|d(\textbf{u},\textbf{u}^{\prime})\leq w;\textbf{u}^{\prime}\in\mathcal{I}\}$, where $w$ is the window size parameter. This parameter determines the size of the square shaped neighbourhood of visible units of side length $2w+1$ pixels, and thus the number of visible variables within the neighbourhood of u is $|\mathcal{N}(\textbf{u})|=(2w+1)^{2}$. These sets of neighbour pixels take the role of the aforementioned subsets $\mathcal{V}(h_{j})$.

However, in order to further control the number of hidden units, we introduce another parameter called stride parameter $t$, which determines how distant each patch of visible units is from each other. The stride parameter, in combination with the window parameter, allows us to control the number of hidden units. These neighbourhoods are distributed such that the centres of each neighbourhood are separated by a minimum distance of $t$. Thus, the set of neighbourhood centres, which is a function of the window and stride parameters $(w,t)$, is defined as $\mathcal{C}(w,t)=\{\textbf{u}|w\leq u_{x},u_{y}\leq d-w\wedge d(\textbf{u},\textbf{u}^{\prime})\geq t;\textbf{u},\textbf{u}^{\prime}\in\mathcal{I}\}$. In order to clarify and remove any ambiguity, we specify that $\textbf{u}=(w,w)$ shall be an element of the set $\mathcal{C}(w,t)$. As can be seen in Figure 1, each hidden unit is only connected to a single subset of pixels and, for illustrative purposes, they are located at the centres given by $\mathcal{C}(w,t)$. Notice that, for our model, there is no need to determine the number of hidden units $n_{h}$, since it is governed by the parameters ($w,t$), where $|\mathcal{C}(w,t)|=n_{h}$. Therefore, we define an SBM model with parameters $(w,t)$ as $\mathcal{M}(w,t)$, where each hidden unit $h_{j}$ is connected only to the visible units $v_{i}\in\mathcal{N}(\textbf{u})$, and there is a one-to-one correspondence between every hidden unit and every $\textbf{u}\in\mathcal{C}(w,t)$.

Nonetheless, there is no need to be restrained to a single set of hidden units determined by $(w,t)$. Several sets of hidden units connected to visible neighbourhoods using different parameters can be stacked, e.g., it is possible to combine two sets of constrained hidden units, where we represent the new stacked model as $\mathcal{M}(w_{1},t_{1};w_{2},t_{2})=\mathcal{M}(w_{1},t_{1})\cup\mathcal{M}(w_{2},t_{2})$.

Such a definition of neighbourhoods may resemble filters from a Convolutional Neural Network (CNN), and although the proposed procedure is inspired by convolutional filters it has nothing to do with them. CNN filters are applied sequentially to patches of pixels in order to obtain a different shaped output, while our neighbourhoods determine connections to a hidden unit (or a range of hidden units) which does not entail a sequential application of a filter.

We compare the performance of our models using benchmark datasets MNIST, FashionMNIST, OrganAMNIST, OrganCMNIST and OrganSMNIST. All of them consist of $28\times 28$ images in grey scale, and even though pixels take discrete values between 0 and 255, they can be easily normalised to the range $[0,1]$, so we can treat the intensities as probabilities. This strategy was used for digit recognition in [30]. Characteristics of each dataset are shown in Table 1.

MNIST consists of images of handwritten digits. It is a benchmark dataset used for the image classification task, and for unsupervised learning with generative models as well. Moreover, in [31], authors contribute by distorting the original MNIST handwritten digits with different corruption types, such that these new images could be used to assess the robustness of classification models against different kinds of noises. This dataset is called the MNIST-C. We use these corrupted images to analyse how well our models perform on the denoising labour as well as on classification. Even though 15 corruptions are available, we use 8 in our experiments. The 7 discarded corruptions involve affine transformation of images, for instance, rotation or scale. It is important that corruptions do not involve the translation, rotation or scaling of the original image, because even though the reconstructed image could be a plausible image for the human eye, the denoising error we use to evaluate execution would be biased and would not be faithful. This may happen because reconstructing an image may remove noisy pixels, but probably will not rotate images to their natural orientation or scale them to the original size. On the other hand, corruptions such as stripe, fog and brightness have high intensity values on the border pixels, which, we have observed, destabilise the RBMs while sampling, since they learn images with border pixels of intensity 0. Therefore, we use corruptions that simply add different types of non-affine noise to the original image or apply a blurring filter, which are the following: shot_noise, impulse_noise, glass_blur, motion_blur, spatter, dotted_line and zigzag.

The FashionMNIST dataset is a set of grey scale images of 10 fashion products, e.g., shoes, shirts and bags. It is commonly used for supervised problems as the original MNIST dataset. Even though neither MNIST nor FashionMNIST have a validation set, during training 10000 instances from the training set are used as the validation set. Thus, the actual size of the train size of MNIST and FashionMNIST is 50000 if we do not state otherwise.

The last three datasets are three of the many available from the MedMNIST dataset. The OrganAMNIST is based on 3D computed tomography (CT) images from Liver Tumor Segmentation Benchmark. Authors used bounding-box annotations of 11 body organs from another study to obtain the organ labels. They also cropped 2D images from the centre slices of the 3D bounding boxes in axial views to obtain the images. Images were resized into 28×28 to perform multi-class classification of 11 body organs. OrganCMNIST and OrganSMNIST were obtained using the same procedure, where the difference comes from cropping the 3D images in Coronal and Sagittal views, respectively. It is worth mentioning that these datasets present unbalanced classes, having two predominant classes. Therefore, when assessing their classification performance, the weighted version of the Log-loss mentioned in Section 2.3 is employed.

Six SBM models defined with different $w,t$ parameters were compared, where, for each SBM, we train another RBM as shown in Table 2, using the same number of hidden units as their related SBM, we refer to them as twin RBM. The parameters that define each model are shown in Table 2. The choice of the $(w,t)$ values was motivated to attain a sufficient number of connections between visible and hidden units so that representational power was not lost, while trying to reduce them as much as possible. Every SBM in Table 2 requires between 6% and 10% of the number of parameters compared to its twin RBM. The SBM and RBM models have been implemented in Python using TensorFlow-2 [35] and the implementation is freely available¹¹1It can be accessed at https://github.com/arkano29/StructuralRBM..

We train the models using mini-batch SGD, using CD with a single Gibbs sampling step for unsupervised learning, while parameters were updated following Eq. (17) on the classification problem. The introduced sparse model can be trained using the learning procedures of the fully connected RBM.

For the sake of increasing the training speed, we apply momentum [36] to the parameter updates. At training epoch $t$, the parameters are updated according to $\Delta\theta(t)=v_{i}(t)=\beta v_{i}(t-1)-\epsilon\frac{\partial E}{\partial\theta}(t)$, where the first term is the accumulated momentum and the second term is the gradient of the objective function with respect to the parameter $\theta$, and $\epsilon$ and $\beta$ are the learning rate and the momentum hyperparameter, respectively. The learning rate is set to $\epsilon=0.1$ and momentum $\beta=0.9$.

We executed the experiments for each model with 10-repeated hold outs, using different seeds, which influences the initialisation of the weight parameters. The weight parameters are initialised randomly following the uniform distribution $U[-\sqrt{6}/\sqrt{n_{v}+n_{h}},\sqrt{6}/\sqrt{n_{v}+n_{h}}]$, as proposed in [37] by Glorot and Bengio, while biases are initialised to 0. As stated by Hinton [38], using a large batch size with mini-batch SGD is a serious mistake, thus he recommends using sizes of the order of tens. Therefore, the batch size is set to 16, and batches of data instances are drawn randomly from the train set at each update step.

Models are trained with 10000 updates of the parameters, and the parameters with the best metric on the validation set during training, which depends on the objective of the training, are chosen for model assessment on the test set. For the unsupervised learning problem, LL is used as the metric to choose parameters, while for the classification problem, the Log-loss is used.

In this section we investigate the performance of SBMs and their twin RBMs for the tasks of probability density estimation, image denoising and image classification.

We train our models to fulfil the unsupervised task of estimating the probability density of the digit and fashion product images introduced in Sec. 5.1. This is fulfilled by training our models to approximate the underlying distribution of the dataset, where the approximation quality is measured using the LL of the model over the test set. The performance of denoising digit images was also assessed. Elapsed time when computing the CD gradient is compared between RBMs and SBMs as well.

Finally, the models have been trained to classify images for different domains, as explained in Sec. 2.2. In addition, their classification robustness against different types of noises have been evaluated for the MNIST images.

RBMs are popular for their unsupervised learning capability and high performance as feature extractors and density estimators [39]. Therefore, the models are trained to approximate the underlying probability distribution of handwritten digits and fashion product images, following the general pipeline described in Section 5.3. Even though both datasets we use provide class labels, they can be considered as unlabelled and thus we train our models for generative purposes using the update rules shown in Eq. (6-8).

Classifying images is one of the most widespread and important problems nowadays, hence we further explore the performance of our models by addressing the image classification problem. As stated in Section 2.2, RBMs can be used as discriminative models. Thus, we can compute the probability of a given image being of a certain class. We have trained our models to classify handwritten digits, fashion products and three organ based medical images from MedMNIST.

Figure 6 presents results of the Log-loss values of the test set after training models with the discriminative update rules in Eq. (17). These experiments highlight the fact that SBMs outperform RBMs in plenty of image domains and problems, thus making them more appropriate to use in classification tasks as well. Fashion product images being the only exception where, for SBMs with less than 265 hidden units, RBMs obtain smaller values of Log-loss values.

Moreover, we highlight the fact that SBM with more than 441 hidden units, having a similar number of trainable parameters compared to RBM₁₂₁, outperform RBM₁₂₁ in all experiments.

It is remarkable that RBM₄₄₁ spikes in Log-loss value when classifying digits compared to every other model, and thus has the worst performance in this specific scenario. Similar to the previous section, Figure 7 shows the Log-loss of the validation set during training for $N_{hidden}\in\{121,441,585\}$, in order to illustrate what could have happened to RBM₄₄₁.

As can be seen in Figure 7, RBM₄₄₁’s Log-loss over the validation decreases slowly when trained with digit images after some epochs, which explains its bad performance. For every other model and image domain, we observe that SBMs always have smaller values of Log-loss compared to their twin RBMs. With the exception of fashion product images, where RBMs and SBMs obtain comparable Log-loss values during training, thus explaining their similar results when assessed with the test set afterwards. Furthermore, the confidence interval illustrates how RBMs have a more unstable training compared to SBMs.