A Large Deviation Theory Analysis on the Implicit Bias of SGD

See Questions.

• Theorems 1,2,4 and proposition 5 consider fixed parameters $\theta$ and study the behavior of the empirical loss under random draws of a dataset $D$. However, in machine learning, the training dataset $D$ is not sampled independent of $\theta$. This is indeed the core challenge of the whole field of learning theory. Thus, I don't understand why Theorems 1,2,4 and proposition 5 are relevant.
• The paper does not actually attempt to provide an full explanation for why SGD leads to better generalization, relative to GD. Instead, the paper merely attempts to shed light on the mechanism by which SGD leads to improved generalization. While this is certainly a valuable goal, I cannot actually follow the paper's logic (see below).
• The crux of the paper is to advocate for decomposition (3) as a meaningful decomposition of the generalization gap. This decomposition decomposes the generalization gap into: (a) the "abnormality" of the training dataset (which measures how unrepresentative the loss on the training dataset is, relative to the loss on randomly sampled datasets of the same size), and (b) the degree to which the loss on randomly sampled datasets of this size is concentrated. Lines 323 - 365 seem to be arguing that the SGD's efficacy can be empirically localized to the first of these factors (the "abnormality"), as opposed to the second (the "concentration"). However, the experiments in Figure 2 show that relative to GD, SGD improves both abnormality and concentration. This seems to contradict the argument in lines 323 - 365 that SGD's success is due to the abnormality alone. This suggests that the decomposition (3) is not a particularly enlightening way of reasoning about the generalization gap.
  • Line 469 seems to be offering an explanation for why SGD also improves the concentration of the loss, but I find the argument convoluted and cannot follow it. To me, the occam's razor explanation for why SGD improves both concentration and abnormality is that this was not an informative decomposition of the generalization gap in the first place.
• Lines 471-486 show that if we have access to the test set, we can use this information to skip certain SGD steps and therefore generalize better. This seems obvious to me, and does not strike me as supporting the paper's theoretical analysis, as claimed on line 482. Further, I would point out that Figure 5 shows that this intervention also improves 'concentration', not just 'abnormality' (which it is intended to do), which further calls into question the utility of this decomposition.

Overall, I'm not convinced that the paper sheds any light on why SGD generalizes better than GD.

Firstly, I have some reservations regarding the suitability of the chosen formulation of Large Deviation Theory (LDT) for addressing the problem of generalization error. Specifically, in the concentration inequalities ((2) and onwards), the probability is considered over the dataset draw while the parameters are held fixed. Informally, it seems likely that the training loss on a given dataset, at the end of training, will typically fall at the far left of the distribution of empirical losses across dataset draws. In some cases, it may be an outlier—a ‘one-off point’ on the left—while the remainder of the distribution might shift significantly to the right and center around the expected loss. Given that LDT is inherently a probabilistic framework aimed at describing the behavior of the tails of distributions, it may not be entirely suited for capturing an outlier that lies apart from the distribution. However, it is precisely the difference between the empirical loss on the training dataset and the expected loss that is of primary interest here. To illustrate, consider an overparameterized neural network setting where training continues until convergence, achieving zero training loss. In these cases, the training data often becomes overfitted—though it can still be beneficial to proceed with training even after the network ‘begins to overfit’ (i.e., when training and validation loss diverge). Once the parameters are fixed, we can evaluate the empirical loss on alternative dataset draws. Assuming the data is sufficiently high-dimensional or exhibits some level of sparsity (such that memorizing the training data alone does not generalize well), evaluating on other dataset draws could approximate the validation or test loss. In scenarios where limited regularization is applied and the model is forced to overfit (for instance, by gradient flow), performance on these 'validation' sets may approach randomness. This would result in a distribution of empirical losses that centers around the expected loss, with the empirical loss on the specific training dataset at zero—thus representing the aforementioned outlier or 'one-off point' within the distribution. However, the actual distributions and the position of the training loss on a given dataset may differ significantly from this assumption. It appears that the issue may stem from the fact that LDT provides bounds for $\hat L(D, \theta)$ rather than $min_\theta \hat L(D, \theta)$. While I could be mistaken, I remain uncertain that this distinction is inconsequential. Given the circumstances described, I am not fully sure that LDT offers meaningful insights into the generalization error. On a somewhat contrasting note, the paper includes experimental examples using datasets with only 50 samples. While I understand this choice may be due to computational constraints, I’m uncertain that it accurately reflects typical deep learning scenarios where models learn specific data representations. Consequently, I’m hesitant to conclude that the distribution graphs presented are truly representative of what occurs in broader deep learning contexts. It might be the case that LDT may serve as a heuristic for decomposing empirical loss in a meaningful way. As mentioned earlier, I think this is a strength of the paper. However, in the specific context of this paper, I am not fully sure that this reformulation offers meaningful insights into the nature of generalization error. For example, simplifying equation (4), the inverse of the rate function approximates a square root (as shown in equation (5)), suggesting that the 'abnormality' of generalization error is approximated by the square of the generalization loss. In this light, it is unclear to me why high abnormality isn’t simply another way of indicating high generalization loss, as stated in the paper; the added focus on abnormality doesn’t immediately seem to provide further benefits beyond analyzing generalization error directly. This concern extends to the observation in the section “SGD PREVENTS HIGHLY ABNORMAL GENERALIZATION ERRORS” (line 402), where it is suggested that smaller batch sizes lead to reduced cosine similarity between the gradient of abnormality for the full dataset and for the batch. If I understand correctly, this point is illustrated solely in Figure 3. With this in mind, it feels though this is a sophisticated reformulation of the generalization error, but lacks sufficient justification for the claim that mini-batch gradient computations lead to smaller generalization errors. Furthermore, if LDT’s role here is largely heuristic, I believe its value would be more convincing with a broader set of experiments.