Why does my regression-NN completely fail to predict some points?

Question

I would like to train a NN in order to approximate an unknown function $y = f(x_1,x_2)$. I have a lot of measurements $y = [y_1,\dots,y_K]$ (with K that could be in the range of 10-100 thousands) which come out from either a simulation or a measurement of a system. I've built a feed-forward NN for solving this problem by using a MSE loss-function, i.e.

$$\mathcal{L} = \frac{1}{K}\sum_{i=1}^K(y_i-\hat{y}_i)^2$$

where I defined as $\hat{y}$ the prediction of the NN. As per activation function I used a ReLU. The network topology is fairly simple with input layer having two neurons $(x_1,x_2)$, three hidden layers with 10 neurons each and finally the output layer (single output neuron).

After the training process expires I've obtained a very curious result. The loss function assumes very small values hence (apparently) indicating that the training is successful. However, if I analyse the squared-error point-by-point, i.e. the quantity

$$ \boldsymbol{\epsilon}_y = [(y_1-\hat{y}_1)^2,\dots,(y_K-\hat{y}_K)^2] $$ I find that for the vast majority of points said quantity is basically zero, with exception of some "outliers" where the error is huge. It looks like this event happens where the gradient of $f()$ is rather big, which might be a reasonable assumption maybe.

I would like for this to not happen anymore. I'd rather accept a slightly bigger error throughout the whole function domain then have the majority of points with null error but some local points that are completely off. As a requirement, the network topology shall be kept rather "easy" so I would not like to increase the number of layers and/or neurons per layer. As a side node, I've also tried to increase the topology complexity a little bit (i.e. 15 neurons per hidden layer and adding a 4th hidden layer) obtaining slighlty better results but still unacceptable error around the steepest points of the function.

I've got two ideas for now:

1. Use a different loss-function $\mathcal{L}$
2. Sample the dataset $(x_1,x_2)$ more frequently around steepest regions, and less frequently where the function is rather smooth and flat

First option

A different loss-function to be adopted. I'm not very familiar with loss-function that might help solve my problem, some rather quick research yielded no good results and I found no significant literature highlighting this kind of problem. I initially thought of something of the likes of

$$ \mathcal{L} = \frac{1}{K}\sum_{i=1}^K(y_i-\hat{y}_i)^2 + \mu\, max\{\boldsymbol{\epsilon}_y\} $$ where $\mu$ is some tuned parameter that penalizes the maximum error. I'm not sure if this makes sense and if it increases the training complexity to the extent that the training process runs in no convergence territory.

Second option

I am not sure how to formalize the undersampling where the function is smooth and flat, I imagine that some kernel used for image processing (edge detection kind of kernels of some sort?) might be helpful but I'm in completely unknown territory.

Conclusion

I'm looking for some insight and ideas (literature results or analogous cases linked are a huge plus!) for solving this curious problems.

Thanks for the help!

Edit

I realize that without having graphical indsights of what I'm doing, it's going to be rather difficult to have a clear understanding. The function $f(x_1,x_2)$ is plotted here below (let alone the red line, it is just there to highlight what happens when $x_1=\text{M}=1.2$ is fixed to a certain value). The missing points of around the variable $x_2=f_n=1$ represent a singularity in my model and I don't care about them being predicted since they physically have no pratical use. I removed them completely for this reason.

Once the training of my NN is complete, I plot the $L2$ error (MSE) between predicted and true value As you can see, it looks like the points where there is a significant gradient change cannot be predicted correctly.

Answer 1

You could try to replace MSE with the recently proposed General and Adaptive Robust Loss Function, a generalization of several losses among which l2 norm (which aside from a square root is equal to MSE). The neat property of this loss is that it uses a dynamic parameter $\alpha$ which makes the loss behaving differently, that can be learned during training. In this way the loss itself adapt to specific outliers in your data (the authors made a video with great visuals to explain this point).

I suspect at first glance (i.e. without data examples) that the problem your dealing with is that steep points are treated by your network as outliers, i.e. your model learn to generalize a smooth function cause the function is smooth most of the time, and surely also cause the amount of parameters of your network is really small (so it has a hard time fitting local irregular patterns). If it true that your function is smooth most of the time than oversampling steep points would be a mistake, cause that would mean training your network on a non representative sample of points, hence if this adaptive loss doesn't improve much your results I would just make the network larger.