Unable to fit linear regression with SGD I'm working on a regression problem, and am unable to fit a single layer feedforward neural network with linear activation via SGD and replicate the performance of sklearn.linear_model. There are 121 features and 9259 training examples, the inputs are normalized to mean=0 and var=1, and I'm using learning_rate=1e-3. As a sanity check, I initialized the network with the coefficients and intercept obtained from sklearn.linear_model and obtained the same MSE=8e-7 and $R^2$=0.51 using both methods. However, the optimization fails completely with random initialization. My data exhibits characteristics that suggest that linear regression may not be the best model, but I'd like to focus this question on replicating the results of sklearn.linear_model.
In order to diagnose the issue, I plotted the following values at each epoch (mean value across minibatches per epoch): weight_norm, prediction, loss, gradient_norm, and update_ratio. update_ratio is defined as learning_rate * gradient_norm/ weight_norm, and I've been monitoring it because Sutskever suggested it should be close to 0.01, which I haven't been able to achieve.
Below is the plot when I initialize the network with the weights obtained from sklearn.linear_model, which immediately obtains $R^2$=0.51 on the training set.

Below is the plot when I initialize the network from the He normal distribution from He et al. (2015), which obtains $R^2$=-21 after 500 epochs, which is significantly worse than just fitting the bias term.

Below is the plot when I initialize with He normal, but scaled down by 0.33 to mirror the weight norm from OLS. This network obtains $R^3$=-3 after 500 epochs.

Additionally, here is the target distribution.

I believe the source of the problem is that the gradient norms are too small, caused by the $\hat{y_i}-y_i$ terms in the gradient being small. Based on how small the gradient norms are, adjusting the learning rate has little effect on training, which I've verified through experimentation.
As a further experiment I've also tried initializing from OLS weights and setting a small proportion of weights to 0 (similar to dropout, but applied to the weights). When I set 3% of weights to 0, the $R^2$ starts off negative and eventually makes it to positive territory, but never to the OLS value of 0.51, even after running for 5000 epochs with various learning rates.
Please let me know what further information I can provide in order to help debug this.
 A: I was able to solve this problem with the following three modifications:


*

*Whitening the features to remove the collinearity

*Normalizing the targets

*Clipping the gradient norms to 1


The network now converges to the OLS loss within a few dozen epochs with learning_rate=1e-3. The optimization fails if I remove any one of the three changes listed above.
I haven't found any literature mentioning the difficulty of training via SGD in the presence of strong collinearity. If anyone is aware of any such sources, please let me know.
Based on comparing the weight norms, SGD appears to have converged to a different set of weights compared to OLS.

A: SGD converges only with step size tending to 0. But I don't think it's the problem here because it matters only for precise convergence, not such big errors. A fixed step size if ok for reasonably precise estimations.
Not sure it is the problem, just trying:
I've observed extremely slow convergence of SGD when the optimal solution overfits a lot. SGD just can't find an overfit solution in reasonable time. This is most often because the optimal weights are "far away" (artificial combinations of big values). In some cases convergence would happen only after millions of epochs while a hundred is enough in normal scenarios.
Regularization solves the problem: the regularized SGD converges as expected. 
The non regularized SGD passes close to the regularized solution after a few steps and then starts desperately looking for the non regularized one practically never converging to it. This fact is due to early stopping being essentially equivalent to (L2) regularization.
