Building an autoencoder in Tensorflow to surpass PCA

Question

Hinton and Salakhutdinov in Reducing the Dimensionality of Data with Neural Networks, Science 2006 proposed a non-linear PCA through the use of a deep autoencoder. I have tried to build and train a PCA autoencoder with Tensorflow several times but I have never been able to obtain better result than linear PCA.

How can I efficiently train an autoencoder?

(Later edit by @amoeba: the original version of this question contained Python Tensorflow code that did not work correctly. One can find it in the edit history.)

Answer 1

Here is the key figure from the 2006 Science paper by Hinton and Salakhutdinov:

It shows dimensionality reduction of the MNIST dataset ($28\times 28$ black and white images of single digits) from the original 784 dimensions to two.

Let's try to reproduce it. I will not be using Tensorflow directly, because it's much easier to use Keras (a higher-level library running on top of Tensorflow) for simple deep learning tasks like this. H&S used $$784\to 1000\to 500\to 250\to 2\to 250\to 500\to 1000\to 784$$ architecture with logistic units, pre-trained with the stack of Restricted Boltzmann Machines. Ten years later, this sounds very old-school. I will use a simpler $$784\to 512\to 128\to 2\to 128\to 512\to 784$$ architecture with exponential linear units without any pre-training. I will use Adam optimizer (a particular implementation of adaptive stochastic gradient descent with momentum).

---

The code is copy-pasted from a Jupyter notebook. In Python 3.6 you need to install matplotlib (for pylab), NumPy, seaborn, TensorFlow and Keras. When running in Python shell, you may need to add plt.show() to show the plots.

Initialization

%matplotlib notebook

import pylab as plt
import numpy as np
import seaborn as sns; sns.set()

import keras
from keras.datasets import mnist
from keras.models import Sequential, Model
from keras.layers import Dense
from keras.optimizers import Adam

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train = x_train.reshape(60000, 784) / 255
x_test = x_test.reshape(10000, 784) / 255

PCA

mu = x_train.mean(axis=0)
U,s,V = np.linalg.svd(x_train - mu, full_matrices=False)
Zpca = np.dot(x_train - mu, V.transpose())

Rpca = np.dot(Zpca[:,:2], V[:2,:]) + mu    # reconstruction
err = np.sum((x_train-Rpca)**2)/Rpca.shape[0]/Rpca.shape[1]
print('PCA reconstruction error with 2 PCs: ' + str(round(err,3)));

This outputs:

PCA reconstruction error with 2 PCs: 0.056

Training the autoencoder

m = Sequential()
m.add(Dense(512,  activation='elu', input_shape=(784,)))
m.add(Dense(128,  activation='elu'))
m.add(Dense(2,    activation='linear', name="bottleneck"))
m.add(Dense(128,  activation='elu'))
m.add(Dense(512,  activation='elu'))
m.add(Dense(784,  activation='sigmoid'))
m.compile(loss='mean_squared_error', optimizer = Adam())
history = m.fit(x_train, x_train, batch_size=128, epochs=5, verbose=1, 
                validation_data=(x_test, x_test))

encoder = Model(m.input, m.get_layer('bottleneck').output)
Zenc = encoder.predict(x_train)  # bottleneck representation
Renc = m.predict(x_train)        # reconstruction

This takes ~35 sec on my work desktop and outputs:

Train on 60000 samples, validate on 10000 samples
Epoch 1/5
60000/60000 [==============================] - 7s - loss: 0.0577 - val_loss: 0.0482
Epoch 2/5
60000/60000 [==============================] - 7s - loss: 0.0464 - val_loss: 0.0448
Epoch 3/5
60000/60000 [==============================] - 7s - loss: 0.0438 - val_loss: 0.0430
Epoch 4/5
60000/60000 [==============================] - 7s - loss: 0.0423 - val_loss: 0.0416
Epoch 5/5
60000/60000 [==============================] - 7s - loss: 0.0412 - val_loss: 0.0407

so you can already see that we surpassed PCA loss after only two training epochs.

(By the way, it is instructive to change all activation functions to activation='linear' and to observe how the loss converges precisely to the PCA loss. That is because linear autoencoder is equivalent to PCA.)

Plotting PCA projection side-by-side with the bottleneck representation

plt.figure(figsize=(8,4))
plt.subplot(121)
plt.title('PCA')
plt.scatter(Zpca[:5000,0], Zpca[:5000,1], c=y_train[:5000], s=8, cmap='tab10')
plt.gca().get_xaxis().set_ticklabels([])
plt.gca().get_yaxis().set_ticklabels([])

plt.subplot(122)
plt.title('Autoencoder')
plt.scatter(Zenc[:5000,0], Zenc[:5000,1], c=y_train[:5000], s=8, cmap='tab10')
plt.gca().get_xaxis().set_ticklabels([])
plt.gca().get_yaxis().set_ticklabels([])

plt.tight_layout()

Reconstructions

And now let's look at the reconstructions (first row - original images, second row - PCA, third row - autoencoder):

plt.figure(figsize=(9,3))
toPlot = (x_train, Rpca, Renc)
for i in range(10):
    for j in range(3):
        ax = plt.subplot(3, 10, 10*j+i+1)
        plt.imshow(toPlot[j][i,:].reshape(28,28), interpolation="nearest", 
                   vmin=0, vmax=1)
        plt.gray()
        ax.get_xaxis().set_visible(False)
        ax.get_yaxis().set_visible(False)

plt.tight_layout()

One can obtain much better results with deeper network, some regularization, and longer training. Experiment. Deep learning is easy!

Answer 2

Huge props to @amoeba for making this great example. I just want to show that the auto-encoder training and reconstruction procedure described in that post can be done also in R with similar ease. The auto-encoder below is setup so it emulates amoeba's example as close as possible - same optimiser and overall architecture. The exact costs are not reproducible due to the TensorFlow back-end not being seeded similarly.

Initialisation

library(keras)
library(rARPACK) # to use SVDS
rm(list=ls())
mnist   = dataset_mnist()
x_train = mnist$train$x
y_train = mnist$train$y
x_test  = mnist$test$x
y_test  = mnist$test$y

# reshape & rescale
dim(x_train) = c(nrow(x_train), 784)
dim(x_test)  = c(nrow(x_test), 784)
x_train = x_train / 255
x_test = x_test / 255

PCA

mus = colMeans(x_train)
x_train_c =  sweep(x_train, 2, mus)
x_test_c =  sweep(x_test, 2, mus)
digitSVDS = svds(x_train_c, k = 2)

ZpcaTEST = x_test_c %*% digitSVDS$v # PCA projection of test data

Autoencoder

model = keras_model_sequential() 
model %>%
  layer_dense(units = 512, activation = 'elu', input_shape = c(784)) %>%  
  layer_dense(units = 128, activation = 'elu') %>%
  layer_dense(units = 2,   activation = 'linear', name = "bottleneck") %>%
  layer_dense(units = 128, activation = 'elu') %>% 
  layer_dense(units = 512, activation = 'elu') %>% 
  layer_dense(units = 784, activation='sigmoid')

model %>% compile(
  loss = loss_mean_squared_error, optimizer = optimizer_adam())

history = model %>% fit(verbose = 2, validation_data = list(x_test, x_test),
                         x_train, x_train, epochs = 5, batch_size = 128)

# Unsurprisingly a 3-year old laptop is slower than a desktop
# Train on 60000 samples, validate on 10000 samples
# Epoch 1/5
#  - 14s - loss: 0.0570 - val_loss: 0.0488
# Epoch 2/5
#  - 15s - loss: 0.0470 - val_loss: 0.0449
# Epoch 3/5
#  - 15s - loss: 0.0439 - val_loss: 0.0426
# Epoch 4/5
#  - 15s - loss: 0.0421 - val_loss: 0.0413
# Epoch 5/5
#  - 14s - loss: 0.0408 - val_loss: 0.0403

# Set the auto-encoder
autoencoder = keras_model(model$input, model$get_layer('bottleneck')$output)
ZencTEST = autoencoder$predict(x_test)  # bottleneck representation  of test data

Plotting PCA projection side-by-side with the bottleneck representation

par(mfrow=c(1,2))
myCols = colorRampPalette(c('green',     'red',  'blue',  'orange', 'steelblue2',
                            'darkgreen', 'cyan', 'black', 'grey',   'magenta') )
plot(ZpcaTEST[1:5000,], col= myCols(10)[(y_test+1)], 
     pch=16, xlab = 'Score 1', ylab = 'Score 2', main = 'PCA' ) 
legend( 'bottomright', col= myCols(10), legend = seq(0,9, by=1), pch = 16 )

plot(ZencTEST[1:5000,], col= myCols(10)[(y_test+1)], 
     pch=16, xlab = 'Score 1', ylab = 'Score 2', main = 'Autoencoder' ) 
legend( 'bottomleft', col= myCols(10), legend = seq(0,9, by=1), pch = 16 )

Reconstructions

We can make the reconstruction of the digits with the usual manner. (Top row are the original digits, middle row the PCA reconstructions and bottom row the autoencoder reconstructions.)

Renc = predict(model, x_test)        # autoencoder reconstruction
Rpca = sweep( ZpcaTEST %*% t(digitSVDS$v), 2, -mus) # PCA reconstruction

dev.off()
par(mfcol=c(3,9), mar = c(1, 1, 0, 0))
myGrays = gray(1:256 / 256)
for(u in seq_len(9) ){
  image( matrix( x_test[u,], 28,28, byrow = TRUE)[,28:1], col = myGrays, 
         xaxt='n', yaxt='n')
  image( matrix( Rpca[u,], 28,28, byrow = TRUE)[,28:1], col = myGrays , 
         xaxt='n', yaxt='n')
  image( matrix( Renc[u,], 28,28, byrow = TRUE)[,28:1], col = myGrays, 
         xaxt='n', yaxt='n')
}

As noted, more epochs and a deeper and/or more smartly trained network will give much better results. For example, the PCA reconstruction error of $k$= 9 is approximately $0.0356$, we can get almost the same error ($0.0359$) from the autoencoder described above, just by increasing the training epochs from 5 to 25. In this use-case, the 2 autoencoder-derived components will provide similar reconstruction error as 9 principal components. Cool!

Answer 3

Here is my jupyter notebook where I try to replicate your result, with the following differences:

• instead of using tensorflow directly, I use it view keras
• leaky relu instead of relu to avoid saturation (i.e. encoded output being 0)
  • this might be a reason for poor performance of AE
• autoencoder input is data scaled to [0,1]
  • I think I read somewhere that autoencoders with relu work best with [0-1] data
  • running my notebook with autoencoders' input being the mean=0, std=1 gave MSE for AE > 0.7 for all dimensionality reductions, so maybe this is one of your problems
• PCA input is kept being data with mean=0 and std=1
  • This may also mean that the MSE result of PCA is not comparable to the MSE result of PCA
  • Maybe I'll just re-run this later with [0-1] data for both PCA and AE
• PCA input is also scaled to [0-1]. PCA works with (mean=0,std=1) data too, but the MSE would be incomparable to AE

My MSE results for PCA from dimensionality reduction of 1 to 6 (where the input has 6 columns) and for AE from dim. red. of 1 to 6 are below:

With PCA input being (mean=0,std=1) while AE input being [0-1] range - 4e-15 : PCA6 - .015 : PCA5 - .0502 : AE5 - .0508 : AE6 - .051 : AE4 - .053 : AE3 - .157 : PCA4 - .258 : AE2 - .259 : PCA3 - .377 : AE1 - .483 : PCA2 - .682 : PCA1

• 9e-15 : PCA6
• .0094 : PCA5
• .0502 : AE5
• .0507 : AE6
• .0514 : AE4
• .0532 : AE3
• .0772 : PCA4
• .1231 : PCA3
• .2588 : AE2
• .2831 : PCA2
• .3773 : AE1
• .3885 : PCA1

Linear PCA with no dimensionality reduction can achieve 9e-15 because it can just push whatever it was unable to fit into the last component.