Stable Coresets via Posterior Sampling: Aligning Induced and Full Loss Landscapes

1. Page 4, Algorithm 1: Wouldn't it be more correct to ensure that the union $V_p$ of would be equal to $S$?

We thank the reviewer for raising this point. It is not necessary in our setting for the union of $V_p$ to exactly equal $S$. For example, under a 10 percent data budget with a batch size of 128 and 100 total batches per epoch, we may set $P=10$ and construct each $V_p$ from $R = 500$ examples, selecting $m = 128$ samples from each for training. In this way, $S$ (the final coreset) is formed without requiring full dataset access. This design avoids the high computational cost and selection bias that may arise when scanning the full dataset, as in methods like Craig.

2. Page 6, Theorem 3.3: What is $\sigma_2$ ?

We thank the reviewer for the question. In Theorem 3.3, $\sigma_2$ refers to the variance of the noise used to perturb the model weights as stated in assumption 4 in appendix A.3, while $\sigma_1$ denotes the variance of the stochastic gradient noise due to data sampling. Unlike $\sigma_1$, which arises inherently from minibatch training, $\sigma_2$ is a tunable parameter under our control and can be adjusted to optimize performance. We will clarify this distinction in the main text and reference the relevant derivations in the appendix for completeness.

3. Supplementary, Page 22: What is w*?

We apologize for the typo. The $w*$ is indeed $\bar{w}$.

4. Supplementary, Page 22, equation 15, last inequality: How was the negative term obtained?

we provide the detailed derivation for it below:

$

\begin{split} f(\bar{w}) \int (f(w) - f(\bar{w}))p(w)dw &= f(\bar{w}) \int (f(w) - f(\bar{w}) -\nabla f(\bar{w})(w-\bar{w}) + \nabla f(\bar{w})(w-\bar{w}))p(w)dw \ &= f(\bar{w}) \int ((f(w) - f(\bar{w}) -\nabla f(\bar{w})(w-\bar{w}))p(w)dw + f(\bar{w})\int \nabla f(\bar{w})(w-\bar{w}))p(w)dw \ &= f(\bar{w}) \int ((f(w) - f(\bar{w}) -\nabla f(\bar{w})(w-\bar{w}))p(w)dw \ &\leq ||f(\bar{w})|| \int ||(f(w) - f(\bar{w}) -\nabla f(\bar{w})(w-\bar{w})||p(w)dw \ &\leq ||f(\bar{w})|| \int c_1||w-\bar{w}||p(w)dw \ & = c_1 ||f(\bar{w})|| E_p[||w-\bar{w}||] \ &\leq c_1 ||f(\bar{w})|| \sqrt{E_p||w-\bar{w}||^2} \ &= c_1 ||f(\bar{w})|| \sigma \sqrt{d} \end{split}

$

The second term at the second line of the derivation is canceled as $\int w p(w) dw = \bar{w}$ and the first inequality is due to the Cauchy-Schwarz inequality and the second inequality is due to the equation 14 in appendix A.1. The final inequality uses Jensen’s inequality due to the concavity of the square root function. Since we obtain the upper bound of the magnitude, we can lower bound the term by switching the sign of the term to be negative as claim in equation 15 and 16 in the appendix A.2. We acknowledge minor issues in the previous version (e.g., a missing constant or square root), but these do not affect the overall dependencies in the result. We will correct these technical details in the updated version.

5. Supplementary, Page 23, equation 20: How was the first term obtained?

we provide the detailed derivation for it below:

$

\begin{split} \int ||f(w) - f(\bar{w}) - \nabla f(\bar{w})(w - \bar{w})||^2 p(w)dw &\geq \int c_2 ||w-\bar{w}||^2p(w)dw \ &=c_2 E_p[||w-\bar{w}||^2] \ &=c_2 \sigma^2d \end{split}

$

The first inequality is due to the equation 14 in appendix A.2.

6. Please ensure to update the proofs and ensure that every pass is justified correctly, since the theoretical contribution of the paper revolves around these proofs.

We thank the reviewer for raising the problem and help make the paper better. We have double checked the proofs and the result remains intact. We will update the proof in the updated version.

7. Why did the authors not compare against GradMatch?

We thank the reviewer for raising this point. The short answer is that GradMatch is already outclassed by other methods we compare against e.g. CREST extensively evaluated against GradMatch and showed the former is superior.

In designing our experimental comparison, we aimed to evaluate a diverse range of methods that reflect fundamentally different principles of subset selection. Specifically, we compared against: (1) Craig, which uses submodular greedy selection based on gradient similarity; (2) Crest, which adaptively selects points to match loss landscape curvature; (3) Glister, which employs bi-level optimization to maximize validation performance; and (4) random selection as a baseline. GradMatch is indeed another gradient matching method, but its core idea—greedy selection based on gradient alignment—is conceptually similar to Craig, with the main difference being the use of Orthogonal Matching Pursuit (OMP) rather than submodular cover. Since Craig already serves as a strong representative of this class, and other methods like Crest have shown strong empirical performance over both, we chose to focus our comparisons on methods with greater conceptual diversity. We will clarify this rationale in the updated version.

8. How does your approach fare against sensitivity-based coreset, e.g., [1]? I am asking since this approach outperformed that of GradMatch.

We thank the reviewer for suggesting another baseline. The sensitivity-based baseline takes a more `global' perspective - constructing a non-uniform distribution to optimize over the worst case bound for an RBFNN on the original data. this in turn acts as a proxy for gradient matching, assuming the gradients can be closely approximated by an RBFNN. However, since the coreset is selected using worst case bounds to be able to fit /any/ RBFNN, it may not capture nuances of data distributions, especially in presence of noise. This is exactly what we observe too, as the following tables show. Note that we were not able to reproduce the results presented in their paper, may be due to different data setups even the random subsampling reported in that paper is way worse than what we observed, and what many other coreset papers have reported. (potentially due to update frequency, learning rate, architecture and other hyperparameters.) Nevertheless we use the code they made publicly available for our evaluations. We see that the method performs better than Craig but worse than crest which is worse than our method. This is especially apparent in the presence of noise.

Cifar-100

Cifar-10

[A] Tukan, M., Zhou, S., Maalouf, A., Rus, D., Braverman, V., & Feldman, D. (2023, July). Provable data subset selection for efficient neural networks training. In International Conference on Machine Learning (pp. 34533-34555). PMLR.

We thank the reviewer for all the comment and address them in the following. If there is still confusion, we will be glad to provide more information or illustration as needed in the discussion stage.

1. Section 2, definition of l: Since the proofs of Theorem A.1 (Appendix A.2) and of Theorem 3.3 (Appendix A.3) require assumptions on loss, assumptions on should be described in the main text. If it is not easy to describe in short, it may be good to provide the assumptions in the appendix and declare in the main text that the detailed assumption is not so restrictive but described in the appendix.

We thank the reviewer for pointing this out. We will explain the assumption in main context or appendix to make it clear.

2. Section 3, Overall: In this algorithm, if we take $\sigma$ to zero does it imply the conventional method? I wondered to what extent the theoretical properties shown in this section holds for the conventional method.

Yes, when $\sigma \to 0$, our method reduces to conventional deterministic approaches such as Craig and Crest. The gradient perturbations vanish, and the selection becomes equivalent to greedy optimization over exact gradients. Importantly, our theoretical results are consistent in this limit. Also, when $\sigma > 0$, our framework introduces loss landscape matching, which helps align coreset selection with the local geometry of the loss landscape. This provides a principled way to improve robustness over standard methods, offering both a generalization and a deeper interpretation.

3. Section 3, Algorithm 1: Can the procedure of taking expectation $E_\delta$ be described explicitly? Is it computed analytically or numerically? In addition, the paper states that "we focus on our attention for sampling in the batch normalization layers" (lines 178-179). How is it implemented in Algorithm 1?

We compute the expectation over $\delta$ numerically by averaging gradients over 4 random perturbations (by sampling $\delta$ with gaussian distribution and add them to the parameter weight) of the model weights. For efficiency and stability, we restrict perturbations to batch normalization layers. This choice significantly reduces memory cost and avoids destabilizing the gradient scale, while still improving performance. Although perturbing all parameters is theoretically appealing, we found that focusing on batch norm layers offers a good trade-off in practice.

4. Section 3, subsection "Algorithm" (from line 158): Although theoretical analysis is done for the problem of equation (6) that finds the smallest coreset whose difference of gradients is $\epsilon$ or less, in reality Algorithm 1 selects a coreset of predetermined size whose difference of gradients is $\epsilon$ smallest. In this sense, can the difference of gradients be represented by M, R and/or T? In other words, can we derive the convergence analysis or other theoretical properties without (only with parameters controllable in Algorithm 1)?

We thank the reviewer for this thoughtful question. The main theoretical tool we rely on is greedy selection for submodular cover problems, where the approximation guarantee is of the form $(1 - \exp(-\gamma))$, assuming $\gamma$-weak submodularity. However, the constant $\gamma$ depends on properties of the loss landscape and the feature geometry, which are data-dependent and not explicitly controlled by $M$, $R$, or $T$ in Algorithm 1. Therefore, this quantity cannot be determined explicitly through the control parameters presented in the algorithm and remain it in symbolic format in our analysis. Empirically, $\epsilon$ is small (see details in Craig paper) with greedy algorithm supporting the practical effectiveness of our fixed-size selection strategy. The theoretical gap between $\epsilon$-based guarantees and fixed-size selection is common in coreset literature (e.g., Craig and Crest), and we adopt a similar pragmatic approach.

5. Section 3, Theorem 3.2: The statement (2) presents the difference between newton step of two subset without dependency of $\sigma$, but is it alright that we ignore in this? If so, why?

Thank you for the insightful question. While the simplified statement in the main text omits $\sigma$ for clarity, the full derivation in Appendix A.2 (e.g., Eq. 33) shows that certain terms do depend on $\sigma$, such as $(\frac{\epsilon}{\sigma})^{1/2}$. However, this dependency is entangled with other factors like the eigenvalues of the Hessian. We chose to highlight the dominant structural terms in the main text but agree it's helpful to clarify $\sigma$'s role. We will revise the paper to make this dependency more explicit.

6. Section 3, Theorem 3.3: In the main text, the effect of $\xi_1$ (random sampling of batches) is not stated anywhere. As far as reading the appendix, does it mean that $\xi_1$ can be computed from $M$, $R$ and/or $T$? (If so, defining $\xi_1$ and $\xi_2$ as the variables with similar functionality is confusing; it should be clarified that is determined from other variables but we need to give assumptions on $\xi_2$.)

We appreciate reviewer for pointing this out. In our analysis, the structure of noise is assumed in equation in assumption part of the proof (see appendix A.2 and equation (49) (50) (57)) (The $\xi_1$ is isotropic noise with variance $\sigma_1$) as done in Crest. We will add additional note for this in the main content to make it clear.

7. Section 4, Overall: In the experiments, $\sigma$ is selected via cross-validation, but what effect $\sigma$ is expected if larger or smaller is selected than the optimal one? Also, it is better if the effect of changing is examined.

In our analysis, we find there exist trade off between matching of gradient and loss landscape. If we increase the magnitude of the noise injected into the model for selection, we will find that there exist better match in loss landscape for current coreset, and decrease of the noise will lead to more accurate estimation of the gradient at current model weight. The balance between the two gives the optimal performance. We also find if we increase the noise by too much, it will start to resemble the random sampling as the gradient becomes too chaotic and cannot capture either the loss landscape or gradient.

8. Section 4, line 288: What is the "design of more effective posteriors"? As far as my understanding, it is not stated in Section 3 or Appendices. What is the mathematical formulation? How can it be implemented?

Thank you for pointing this out. The "design of more effective posteriors" refers to the flexibility in choosing how we perturb the model weights during gradient estimation. The current strategy can be interpreted as drawing samples from an isotropic Gaussian posterior centered at $\bar{w}$ with variance $\sigma^2$ (i.e., uniform in all directions). However, one could consider more informed posterior distributions. For instance, perturbing only along sharpest directions in the weight space—akin to using an anisotropic or structured posterior—may better explore sensitive regions of the loss landscape. While Section 4 outlines this idea at a high level, implementation details for our current perturbation strategy are provided in Appendix A.4. We agree that further formalization of alternative posteriors is an exciting direction for future work and will clarify this in the revision.

9. Section 2, lines 121-122: What does the sentence "For tractability and better generalization, in practice, we use SGD (stochastic gradient descent) for optimization of (1) can lead to large memory burden and can be inefficient due to the calculation time" intend? As long as I read this sentence, I understood that "we usually use SGD but it is inefficient", but it is somewhat unnatural.

We appreciate that the reviewer pointing this out. There should be gradient descent (GD) in the front for comparison.

We appricate the reviewer comment and address them in the following:

1. Can you provide results on one more language or multimodel dataset on small models? It would make the result more concrete.

We thank the reviewer for the suggestion. To strengthen our claims, we conducted an additional experiment on the TREC-50 language dataset using ELECTRA-small-discriminator, a compact transformer model with approximately 14 million parameters. Alongside SNLI (used in the main paper), this dataset represents a language inference-style task. We finetune the pretrained model for 50 epochs under a 10 percent data budget, using AdamW with a learning rate of 1e-4, weight decay of 0.01, and standard default settings. We also apply learning rate warm-up for 1 epoch to stabilize training. The results below show that our method outperforms both Crest and random selection under various corruption level:

2. Can you theoretically compute time complexity of Crest, Craig and your method? (I have seen your method's complexity on appendix C but also want to see the difference and comparison of others.) If possible, please provide system runtime for training as well.

We thank the reviewer for asking the question. The main analysis of time complexity lies on the selection of subset and we provide comprehensive comparison here.

For Craig, there exist the need for forward propogation for each sample, and the corresponding time complexity is $\mathcal{O}(dn)$. (d:model parameter, n:size of the dataset) The method is to select $q$ fraction ($0 < q < 1$) of the whole dataset with size all at once. The greedy selection strategy for selecting $qn$ samples out n samples is $\mathcal{O}(qn^2)$. Therefore, the complete time complexity is $\mathcal{O}(qn^2 + dn)$ for each epoch.

For Crest, instead of selecting from whole dataset, they select from subset with size M, and the corresponding time complexity for forward propogation is $\mathcal{O}(dR)$. (d:model parameter, n:size of the dataset) The method is to select $m$ point from the subset. The greedy selection strategy for selecting $Rm$ samples out n samples is $\mathcal{O}(Rm)$. Therefore, the time complexity is $\mathcal{O}( P (Rm+ dR))$ where the $P$ is the number of subset that is manually chosen. The overall time complexity for Crest method need to multiple by the times they update the coreset within one epoch for fair comparison, but the average number of update for coreset cannot be calculated as they utilize adaptive strategy which depend on the curveture of the loss landscape.

For our method, we also select from subset with size R but we need multiple forward propogation (M times) and the corresponding time complexity is $\mathcal{O}(MdR)$. We then select $m$ points from the subset and the corresponding time complexity is $\mathcal{O}(mR)$. We need to perform multiple time to have $q$ portion of the whole dataset. The overall time complexity is $\mathcal{O}(q\frac{n}{R} (MdR+mR))$ for one epoch.

3. The paper does not report actual training time and only provides theoretical and data usage analysis. I would encourage the authors to include a more comprehensive evaluation of system performance.

We appreciate the reviewer’s suggestion. We acknowledge the importance of reporting practical system performance in addition to theoretical analysis. We have already included some relevant metrics: Memory usage across the entire training process is presented in Figure 2. Training time for different datasets is provided in Appendix A.5, Table 5. A brief comparison of training time between Crest and our method is discussed in the caption of Table 1 in the main paper. In the updated version, we will include a dedicated table that summarizes training time, memory usage, and coreset selection overhead across methods for a more comprehensive system-level evaluation.

4. This work has not been tested on larger models and datasets, which weakens its generality. The paper does not report actual training time and only provides theoretical and data usage analysis. I would encourage the authors to include a more comprehensive evaluation of system performance.

We thank the reviewer for raising this concern. To address the issue of generality, we would like to clarify that we have conducted extensive evaluations across a wide range of datasets and model sizes: Datasets: Our experiments include MNIST, CIFAR-10, CIFAR-100, TinyImageNet, ImageNet, and SNLI (a language inference dataset), with dataset sizes ranging from 60,000 to 1.4 million samples. Model scales: We evaluate models from small to large, including LeNet, ResNet-20, ResNet-18, ResNet-50, RoBERTa, and ViT-Base, spanning parameter counts from 10K to over 100M. (See Appendix A.9 for ViT results.) Robustness checks: We test under varying data corruption ratios (0 to 90 percent), data budgets (1, 10, 20 percent), and 3–5 random seeds to ensure consistent performance. As for system performance, we report memory usage in Figure 2 and include training time comparison in Appendix A.5, Table 5. A more detailed runtime and memory comparison table will be added in the updated version for completeness. To our knowledge, little prior work on coreset selection has systematically evaluated across such a broad combination of model sizes, dataset types, corruption levels, and data budgets. We believe this provides a comprehensive demonstration of the method’s robustness and scalability.

We appreciate Reviewer for the comment and summarize the quesions and responses in the following:

1. Have any experiments been performed on transformer based models?

Yes. As shown in Appendix A.9, Table 8, we conduct experiments using ViT-Base on the CIFAR-100 dataset, and our method consistently outperforms other baselines across different corruption ratios. Additionally, we evaluate our method on transformer-based language models, specifically RoBERTa (a variant of BERT utilizing attention-based mechanisms), on the SNLI natural language inference dataset. (in the main context) In both the vision and language domains, our approach demonstrates improved performance over standard baselines, suggesting its general applicability to transformer architectures.