Effect of rescaling of inputs on loss for a simple neural network

Question

I've been trying out a simple neural network on the fashion_mnist dataset using keras. Regarding normalization, I've watched this video explaining why it's necessary to normalize input features, but the explanation covers the case when input features have different scales. The logic is, say there are only two features - then if the range of one of them is much larger than that of the other, the gradient descent steps will stagger along slowly towards the minimum.

Now I'm doing a different course on implementing neural networks and am currently studying the following example - the input features are pixel values ranging from 0 to 255, the total number of features (pixels) is 576 and we're supposed to classify images into one of ten classes. Here's the code:

import tensorflow as tf

(Xtrain, ytrain) ,  (Xtest, ytest) = tf.keras.datasets.fashion_mnist.load_data()

Xtrain_norm = Xtrain.copy()/255.0
Xtest_norm = Xtest.copy()/255.0

model = tf.keras.models.Sequential([tf.keras.layers.Flatten(),
                                    tf.keras.layers.Dense(128, activation="relu"),
                                    tf.keras.layers.Dense(10, activation="softmax")])

model.compile(optimizer = "adam", loss = "sparse_categorical_crossentropy")
model.fit(Xtrain_norm, ytrain, epochs=5)
model.evaluate(Xtest_norm, ytest)
------------------------------------OUTPUT------------------------------------
Epoch 1/5
60000/60000 [==============================] - 9s 145us/sample - loss: 0.5012
Epoch 2/5
60000/60000 [==============================] - 7s 123us/sample - loss: 0.3798
Epoch 3/5
60000/60000 [==============================] - 7s 123us/sample - loss: 0.3412
Epoch 4/5
60000/60000 [==============================] - 7s 123us/sample - loss: 0.3182
Epoch 5/5
60000/60000 [==============================] - 7s 124us/sample - loss: 0.2966
10000/10000 [==============================] - 1s 109us/sample - loss: 0.3385
0.3384787309527397

So far, so good. Note that, as advised in the course, I've rescaled all inputs by dividing by 255. Next, I ran without any rescaling:

import tensorflow as tf

(Xtrain, ytrain) ,  (Xtest, ytest) = tf.keras.datasets.fashion_mnist.load_data()

model2 = tf.keras.models.Sequential([tf.keras.layers.Flatten(),
                                    tf.keras.layers.Dense(128, activation="relu"),
                                    tf.keras.layers.Dense(10, activation="softmax")])

model2.compile(optimizer = "adam", loss = "sparse_categorical_crossentropy")
model2.fit(Xtrain, ytrain, epochs=5)
model2.evaluate(Xtest, ytest)
------------------------------------OUTPUT------------------------------------
Epoch 1/5
60000/60000 [==============================] - 9s 158us/sample - loss: 13.0456
Epoch 2/5
60000/60000 [==============================] - 8s 137us/sample - loss: 13.0127
Epoch 3/5
60000/60000 [==============================] - 8s 140us/sample - loss: 12.9553
Epoch 4/5
60000/60000 [==============================] - 9s 144us/sample - loss: 12.9172
Epoch 5/5
60000/60000 [==============================] - 9s 142us/sample - loss: 12.9154
10000/10000 [==============================] - 1s 121us/sample - loss: 12.9235
12.923488986206054

So somehow rescaling does make a difference? Does that mean if I further reduce the scale, the performance will improve? Worth trying out:

import tensorflow as tf

(Xtrain, ytrain) ,  (Xtest, ytest) = tf.keras.datasets.fashion_mnist.load_data()

Xtrain_norm = Xtrain.copy()/1000.0
Xtest_norm = Xtest.copy()/1000.0

model3 = tf.keras.models.Sequential([tf.keras.layers.Flatten(),
                                    tf.keras.layers.Dense(128, activation="relu"),
                                    tf.keras.layers.Dense(10, activation="softmax")])

model3.compile(optimizer = "adam", loss = "sparse_categorical_crossentropy")
model3.fit(Xtrain_norm, ytrain, epochs=5)
model3.evaluate(Xtest_norm, ytest)
------------------------------------OUTPUT------------------------------------
Epoch 1/5
60000/60000 [==============================] - 9s 158us/sample - loss: 0.5428
Epoch 2/5
60000/60000 [==============================] - 9s 147us/sample - loss: 0.4010
Epoch 3/5
60000/60000 [==============================] - 8s 141us/sample - loss: 0.3587
Epoch 4/5
60000/60000 [==============================] - 9s 144us/sample - loss: 0.3322
Epoch 5/5
60000/60000 [==============================] - 8s 138us/sample - loss: 0.3120
10000/10000 [==============================] - 1s 133us/sample - loss: 0.3718
0.37176641924381254

Nope. I divided by 1000 this time and the performance seems worse than the first model. So I have a few questions:

1. Why is it necessary to rescale? I understand rescaling when different features are of different scales - that will lead to a skewed surface of the cost function in parameter space. And even then, as I understand from the linked video, the problem has to do with slow learning (convergence) and not high loss/inaccuracy. In this case, ALL the input features had the same scale. I'd assume the model would automatically adjust the scale of the weights and there would be no adverse effect on the loss. So why is the loss so high for the non-scaled case?

2. If the answer has anything to do with the magnitude of the inputs, why does further scaling down of the inputs lead to worse performance?

Does any of this have anything to do with the nature of the sparse categorical crossentropy loss, or the ReLU activation function? I'm very confused.

Answer 1

Modern initialization methods are designed with strong assumptions about the scale of the input data, usually inputs have 0 mean and unit variance or that inputs are in the unit interval.

If we apply scaling so that inputs are $X_{ij}\in [0,1]$, then activations for the first layer during the first iteration are $$ X\theta^{(1)} + \beta^{(1)} $$

and at convergence are $$ X\theta^{(n)} + \beta^{(n)}, $$ where the weights are $\theta$, the bias is $\beta$.

Network initialization draws values from some specific distribution, usually concentrated in a narrow interval around 0. If you don't apply scaling, then activations for the first layer during the first iteration are $$ 255\cdot X\theta^{(1)} + \beta^{(1)} $$ So the effect of multiplying by the weights is obviously 255 times as large. At convergence, the model will arrive at the same loss as the scaled case; however, it will take longer to get there since the non-scaled model has initial weights that are 255 times too large. Larger weights are close to saturating the softmax function, where the gradient is not very steep, so it will take a long time for the weights to update enough to compensate for the lack of scale.

---

The video is talking about scaling in the context of features that have different scales. That poses a different kind of conditioning problem for the optimizer. See: In Machine learning, how does normalization help in convergence of gradient descent?