Poor convergence of a neural network, which implements NMF

Question

I'd like to understand why this simple network fails to converge. The resulting MSE error is an order of 10^4 - 10^5 bigger than what could be achieved. The task is to do a non-negative matrix factorization when $\mathbf{V}$ is partially observed. For example we know the average weekly sales of goods in grocery stores, but of course not all products are sold in all stores. This example ignores seasonality, trends etc.

First I generate random data, and do a 50-50 train-test split:

import numpy as np
import matplotlib.pyplot as plt


n_prod, n_store, rank = 1000, 100, 12

weight = 10.0**np.linspace(0, -2, rank)  # Each extra factor has less impact on average

Ap = np.exp(np.random.randn(n_prod,  rank)) * weight[None]
As = np.exp(np.random.randn(n_store, rank)) * weight[None]

sales = Ap.dot(As.T)
assert sales.shape == (n_prod, n_store)


plt.clf()
plt.subplot(1,2,1)
plt.plot(Ap.std(axis=0)); plt.plot(As.std(axis=0)); plt.grid(True)
plt.legend(['Product factors', 'Store factors']); plt.xlabel('Factor index'); plt.ylabel('STD')

plt.subplot(1,2,2)
plt.hist(np.clip(sales.flatten(), 0, np.percentile(sales, 98)), 100); plt.grid(True)
plt.ylabel('Product-store sales')
plt.tight_layout()


seen_data = np.random.random((n_prod, n_store)) < 0.5


# assert that each product and store has sufficent training data
assert seen_data.sum(axis=0).min() > np.sqrt(n_prod)
assert seen_data.sum(axis=1).min() > np.sqrt(n_store)

Then I train a network 10 times and plot losses:

from tensorflow import keras
from tensorflow.keras import layers as l
import tensorflow.keras.backend as K


model_rank = 8

err = sales - Ap[:,:model_rank].dot(As[:,:model_rank].T)
ref_mse = (err[~seen_data]**2).mean()
print(f"Theoretical MSE: {ref_mse}")
# Theoretical MSE: 7.160030917701594e-05

ix_prod, ix_store = np.where(seen_data)
Y = sales[seen_data]

ix_prod_val, ix_store_val = np.where(~seen_data)
Y_val = sales[~seen_data]


hs = []  # Collect histories here
verbose = 0

for _ in range(10):
    inputs = (l.Input(1), l.Input(1))
    prod_emb  = l.Embedding(n_prod,  model_rank, embeddings_regularizer=keras.regularizers.l2(1e-6))
    store_emb = l.Embedding(n_store, model_rank, embeddings_regularizer=keras.regularizers.l2(1e-6))
    
    # Embeddings are in log-scale
    x_prod  = K.exp(prod_emb(inputs[0]))
    x_store = K.exp(store_emb(inputs[1]))
    
    # Dot-product between factors
    fc = K.sum(x_prod * x_store, axis=-1)
    
    model = keras.models.Model(inputs, fc)
    model.summary()
    
    
    model.compile(optimizer=keras.optimizers.Adam(learning_rate=1e-3), loss='mse')
    
    h = model.fit([ix_prod, ix_store], Y, batch_size=2**13, epochs=10000, verbose=verbose,
                  validation_data=((ix_prod_val, ix_store_val), Y_val),
                  callbacks=[keras.callbacks.EarlyStopping(monitor='loss', patience=50, restore_best_weights=True, min_delta=0.01),
                             keras.callbacks.ReduceLROnPlateau(patience=20, min_delta=0.1, min_lr=1e-4, verbose=verbose)])
    hs.append(h.history)

print(f"Number of epochs: {[len(h['loss']) for h in hs]}")
# Number of epochs: [1118, 1095, 1182, 1149, 1044, 1079, 1020, 1102, 1124, 1008]

# Don't plot first epochs, we are interested on the final convergence.
n_skip = 50

plt.clf()
plt.subplot(1,2,1); [plt.plot(h['loss'][n_skip:]) for h in hs]; plt.grid(True); plt.ylim(0)
plt.title(f"loss, min={min(l for h in hs for l in h['loss'])}")

plt.subplot(1,2,2); [plt.plot(h['val_loss'][n_skip:]) for h in hs]; plt.grid(True); plt.ylim(0)
plt.title(f"val_loss, min={min(l for h in hs for l in h['val_loss'])}")
plt.suptitle(f"Reference MSE for {model_rank} out of {rank} factors = {ref_mse}")
plt.tight_layout()

Results are quite bad, and the network stops converging early.

For reference, I used a closed-source NFM code (which actually minimizes KL-divergence and not MSE) and its fit MSE is about 0.1 and validation MSE is about 0.3, or more than 10x less than with Keras. Granted, it is still quite far from the theoretical error which was calculated by using the same factors which were used to generate the data.

Sometimes I got acceptable results with SGD with momentum, but the results were very sensitive on momentum and learning rate hyperparameters. Here are several plots:

Actually sometimes with high learning rates Adam and its variants manage to overcome the stagnation. But I haven't seen this kind of behavior on other neural network applications. So, any guesses what is happening here?

Edit:

Lpounng pointed out that the problem is non-convex, and thus difficult to optimized. But the closed-source implementation finds very good solution nevertheless. So maybe this is more of an issue with general purpose algorithms.

Anyway, this reminded me that neural networks seem to escape local minima due to having saddle-points in the high-dimensional parameter space (citation needed). So I tried changing the network structure so that embeddings aren't summed or multiplied together but rather just concatenated. These are then squeezed down to the one output dimension in a varying number of steps. Adam optimizer starts from a learning rate of 3e-3, and I added the final LR to the legend to check whether it reaches the configured minimum of 1e-4. Used dimensions from 20 to 1 are shown in plot titles. np stands for "number of trainable parameters", which doesn't increase much.

These results are already much better, and the MSE loss is about 0.5 - 1. Here intermediate dimensions were interpolated in log-scale. I also tried a linear scale. I left out the case when n_steps is zero, as it is identical to the previous structure.

So it seems that the original network structure was just ill-suited for general purpose optimizers. Source code to generate model variants:

inputs = (l.Input(1), l.Input(1))
prod_emb  = l.Embedding(n_prod,  model_rank, embeddings_regularizer=keras.regularizers.l2(1e-6))
store_emb = l.Embedding(n_store, model_rank, embeddings_regularizer=keras.regularizers.l2(1e-6))

x = K.concatenate([prod_emb(inputs[0]), store_emb(inputs[1])])

# Add two to `n_steps`, since we'll trim out the first and last values. Choose one of these.
dims = np.exp(np.linspace(np.log(x.shape[-1]), 0, n_steps+2)).round().astype(int)
dims = np.linspace(x.shape[-1], 1, n_steps+2).round().astype(int)

if n_steps > 0:
    # Make `dims` a descending list, so that it doesn't repeat any values.
    for i in range(len(dims)-1, 0, -1):
        if dims[i-1] <= dims[i]:
            dims[i-1] = dims[i] + 1
    
    n_steps2dims[n_steps] = dims
    
    # Skip the first and last element
    dims = dims[1:-1]
    
    for dim in dims:
        x = l.BatchNormalization()(l.Dense(dim, activation='elu')(x))

x = l.Dense(1, activation='exponential')(x)

Answer 1

Haven't gone through the whole code, but NMF is always non-convex, thus we may end up in a local minima given some initial condition and optimization setting. This also aligns with your observation that sometimes high learning rate helps jump out of the local minima trap.

In fact this kind of behavior is very common in neural network.