Custom Loss based on Input Features

I built a neural network to solve a regression problem for me. I would like to wrap a physical equation around it which depends on the input features. Instead of predicting Y directly, I want to predict a parameter alpha which is part of a physical equation.

1. Let the model predict alpha
2. Instead of calculating the loss as MSE(y_true, y_pred) directly, calculate y_pred as follows before (simple example): y_pred = alpha * X[0] + X[1] where X is the vector containing the model input, alpha the parameter the NN should learn.

This is how I implemented it, you can see the custom_loss function and part where I call the fit() method. This is a workflow for testing I found on stackoverflow, if you know a better method to implement something like this, great!

def custom_loss(X, alpha): # Get target observations from data y_true = X[:, 0] num = X[:2, 0] y_pred = (alpha * X[0] + X[1] ) mse = tf.keras.losses.MeanSquaredError() return mse(y_true, y_pred) model.fit(train_data["Xtrain"], np.append(train_data["Ytrain"], train_data["Xtrain"], axis=1), epochs=10000, batch_size=200, validation_data=(train_data["Xval"], np.append(train_data["Yval"], train_data["Xval"], axis=1))) 

However, my model is super unstable and doesn’t converge. Sometimes I get only NaNs as loss, sometimes the loss circles around a value but doesn’t finish the thing. Using only MSE, the model works quite well. Let’s assume the physical equation embedded in the loss function “makes sense”, this is just an example.

Did anybody do something like this before? Another problem I am facing – the input X is normalized (StandardScaler). For the physical equation to make sense, I would need the “de-normalized” values, right? How would I attempt this? scikit-learn works with numpy not with tensors I guess.

Thanks a lot!