You are right that the mean (as well as the median of nonzero elements) in the gradient diff remains small (=O(1e-8)), independent of the dimensionality. Unfortunately, roughly 30-40% of the elements in the gradient vector are different in the two approaches, so a lot of parameters start getting slightly different updates as training progresses.
In my case, the cached version gets consistently worse accuracies than the non-cached version (67.5% vs 69% on my dev set).
I did some more digging of the toy example at dimensionality=1.
As in the toy example, I took Loss = wx + (wx)^2.
So the gradient, dL/dw = x + 2wx^2
I set:
w = 0.8032760620117188
x = 0.17483338713645935
With caching, I get gradient = 0.2239404171705246 from pytorch
Without caching, I get gradient = 0.2239404022693634 from pytorch
And if I compute the gradient myself in the python interpreter, I get gradient = 0.22394040524488334
So all of them are different, and within 1e-8 of each other.
If this is indeed to do with numerical precision issues, these issues cascade as training goes by.
As I mentioned earlier, the discrepancy can be 1e-6 for some elements in the gradient, which take the corresponding parameter along a significantly different optimization path.
I am not sure what the way ahead is at this point. I would sure love to use caching but the loss in accuracy is quite significant.