Sample-Adaptive Zeroth-Order Post-Training
summary
The gist
The gist The authors propose Sample-Adaptive Zeroth-Order Optimization (SAZO), which learns sample importance online and adaptively allocates queries to accelerate convergence.
In short
Sample-Adaptive Zeroth-Order Optimization (SAZO) improves memory-efficient fine-tuning of large language models by learning sample importance online. It estimates how important each training example is using existing directional differences, rather than needing true gradients. This adaptive sampling strategy accelerates convergence, achieving up to 6x faster training with the same query count and better downstream performance.
Key concepts
- Zeroth-Order (ZO) Optimization
- A method used to fine-tune large language models that estimates updates using only forward passes (queries). It is memory-efficient because it avoids the high memory cost of storing full gradients, but standard ZO methods sample examples uniformly, which wastes computational effort.
- Sample Importance Learning
- SAZO learns which training examples are most valuable to update the model online. It does this by analyzing existing directional differences—the results of previous queries—to estimate the importance of each sample without needing access to true gradients or new model evaluations.
- Exponential Moving Averages (EMA)
- These are used in SAZO to track how the importance of a specific training example changes over time. By maintaining per-example EMAs, the method aggregates sparse feedback from minibatch updates to create an evolving picture of sample relevance.
Terminology used across episodes
This episode discusses
- Sample-Adaptive Zeroth-Order Post-Training · Paper Radio
- Stochastic Optimization with Optimal Importance Sampling
- Online Learning to Sample
- FZOO: Fast Zeroth-Order Optimizer for Fine-Tuning Large Language Models towards Adam-Scale Speed
- On Adaptivity in Zeroth-Order Optimization
- Variance-reduced Zeroth-Order Methods for Fine-Tuning Language Models
- LoRA: Low-Rank Adaptation of Large Language Models
- Adam: A Method for Stochastic Optimization
- PRAC: Principal-Random Subspace for LLM Activation Compression and Memory-Efficient Training
- Adam with Bandit Sampling for Deep Learning
- RoBERTa: A Robustly Optimized BERT Pretraining Approach
- Fine-Tuning Language Models with Just Forward Passes
- Stochastic Optimization with Bandit Sampling
- Enhanced Federated Optimization: Adaptive Unbiased Client Sampling with Reduced Variance
- OPT: Open Pre-trained Transformer Language Models
- Second-Order Fine-Tuning without Pain for LLMs:A Hessian Informed Zeroth-Order Optimizer
The paper
Sample-Adaptive Zeroth-Order Post-Training · Read on arXiv
Haiyuan Wan, Yanyi Li, Weicheng Lin, *Cong Fang
School of Physics, Peking University · School of Intelligence Science and Technology, Peking University
Zeroth-order (ZO) optimization offers a memory-efficient alternative to backpropagation for fine-tuning large language models. However, existing ZO methods typically sample training examples uniformly, overlooking their evolving importance and leaving avoidable estimation variance. To address this limitation, we propose Sample-Adaptive Zeroth-Order Optimization (SAZO), which learns sample importance online and adaptively allocates queries to accelerate convergence. SAZO extracts gradient-norm information from existing directional differences to estimate sample importance without access to true gradients. Furthermore, to track evolving importance from sparse minibatch feedback, SAZO maintains per-example exponential moving averages that aggregate observations across updates. The method requires no additional model evaluations and only O(N) additional scalar state compared with MeZO. We prove that SAZO asymptotically matches the leading convergence term of the minimum-variance sampling oracle, and achieves a strictly faster convergence rate than standard methods. Extensive experiments across four model families demonstrate improved average downstream performance and up to 6 times faster convergence to the same training loss in terms of query count.
Transcript
Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.
Tom: Today's paper: "Sample-Adaptive Zeroth-Order Post-Training".
Jane: The gist The authors propose Sample-Adaptive Zeroth-Order Optimization (SAZO), which learns sample importance online and adaptively allocates queries to accelerate convergence.
Tom: First, who's behind it and why it matters.
Title and authors: Tom: So we're talking about this paper today, "Sample-Adaptive Zeroth-Order Post-Training," which is tackling how to fine-tune those huge language models without needing the full backpropagation process.
Jane: Exactly. The authors, Haiyuan Wan, Yanyi Li, Weicheng Lin, and Cong Fang, they propose something called Sample-Adaptive Zeroth-Order Optimization or SAZO. It sounds complicated but the core idea is learning sample importance online to decide which queries to use next for better results.
Lu: It’s about solving a problem where existing zeroth-order methods sample training examples uniformly, meaning they treat every example the same way, which ignores how important certain examples are as training progresses >
Meng: So, if you're running a big model and you only have two forward evaluations to estimate an update, uniform sampling means you’re wasting those evaluations on things that don't matter much right now >
Tom: Right. And the paper says SAZO learns this sample importance adaptively. It extracts information about which samples are important just from the existing directional differences, without needing the true gradients >
Jane: That's a big thing because it bypasses a major hurdle: they don't need any *additional* model evaluations to figure out where to focus their queries next >
Lu: And they track this evolving importance by keeping per-example exponential moving averages, which smooths out the sparse feedback from the minibatch updates across time >
Meng: It sounds like they are trying to make a system that's smarter about its sampling strategy based on what it actually sees happening during training >
Tom: And they do this using importance sampling, where the update is weighted by these estimated norms, specifically using gbt = sIt / N pIt for the update and z t t+one = x t - eta gbt for the next state > <ref:2610.12566#pg1>
Jane: That weighting is key because it aims to match what would be the ideal allocation, which they define as being proportional to the current per-example gradient norms, p* i, which are related to grad f i(x t) >
Lu: They estimate those gradient norms using the scaled magnitudes of existing directional differences, Y i,t = r / pi squared s i,t, and they show that this estimation is accurate up to an error term EtY i,t = a i,t + O(mu two) > <ref:2610.12566#pg1>
Meng: So what's the practical implication for deployment? The paper claims this method requires no additional model evaluations and only O(N) additional scalar state compared to existing methods like MeZO >
Tom: That memory footprint is important. They prove that SAZO asymptotically matches the leading convergence term of the minimum-variance sampling oracle, meaning it gets as good as it can get in terms of variance minimization >
Jane: But they also show it achieves a strictly faster convergence rate than standard methods, and they did extensive experiments across four model families: RoBERTa, OPT-1 point 3B, Phi-two and Llama3-8B >
Lu: The experimental results are pretty strong. They demonstrate improved average downstream performance while maintaining MeZO’s memory footprint and per-update cost >
Title and authors: Meng: And the speedup is what caught my attention most—they achieve up to six times faster convergence to the same training loss using fewer queries than MeZO >
Tom: Six times faster means if you're optimizing a model, you can get there much quicker in terms of how many forward passes you need >
Jane: That’s a significant win for anyone who is doing large scale post-training because it directly translates to reduced computational cost for the optimization phase >
Lu: The analysis on robustness is also interesting. They found that SAZO performs best with a smoothing coefficient beta of zero point zero five, and they show that the gap between uniform and optimal horizon-averaged second moments, U T - D T, is always greater than or equal to zero > <ref:2610.12566#pg2>
Meng: What does that inequality mean for us in practice? It means even with the noise and drift inherent in this estimation process, we are still doing better than just guessing a fixed sampling plan >
Tom: The analysis also breaks down the sources of error, showing that the radius-independent noise and exploration contributions have scales like beta A T epsilon squared, N squared d eta squared D E T epsilon cubed beta squared, and epsilon D E T > <ref:2610.12566#pg2>
Jane: And Theorem three gives a bound on the regret, showing that Regret T / T* three hundred eighty-four K e T/ten + three hundred eighty-four R mu T >
Lu: That bound is what makes it concrete. It shows the performance gains are preserved across different model scales and different training population sizes, which is a big plus for generalization >
Meng: For me, deployment-wise, as long as we can maintain that O(N) additional scalar state requirement, this method seems viable because it keeps the per-update latency low compared to what we're used to with MeZO >
Tom: It really does. So to wrap up on this paper by Haiyuan Wan et al., Sample-Adaptive Zeroth-Order Post-Training, SAZO is a technique that learns sample importance online using existing directional differences without needing true gradients >
Jane: It addresses the problem of uniform sampling by tracking evolving importance through per-example exponential moving averages >
Lu: The main results show it matches the minimum-variance oracle's convergence term asymptotically and achieves up to six times faster query count speedup over MeZO on four different model families >
Meng: It maintains MeZO’s memory footprint while being much faster in terms of the number of queries it needs to reach a target loss >
Tom: It's a method that leverages sample heterogeneity to reduce estimation variance, and based on the evidence, it looks like it performs very well across different model sizes >
Jane: So for anyone listening who doesn't read papers, SAZO means we can use existing directional data to make much smarter decisions about which parts of the training set to query next, leading to faster fine-tuning >
The paper's summary: Tom: So, we're looking at how they wrap up this paper on Sample-Adaptive Zeroth-Order Post-Training, SAZO, which is essentially a system that learns where to focus its queries online using existing data differences >
Jane: Right. We talked about how it tracks sample importance using those moving averages across updates and how it uses them to weight the next query allocation >
Tom: Yeah. The main thing they prove is that SAZO asymptotically matches the leading convergence term of the minimum-variance sampling oracle, which means mathematically, it's as good as it can be in terms of minimizing variance >
Jane: That’s a big theoretical win because it shows their adaptive sampling isn't just a clever trick; it matches the best possible mathematical bound for this type of estimation >
Lu: What I find interesting is how they break down the noise. They show that the radius-independent noise and exploration contributions scale with things like beta A T epsilon squared and N squared d eta squared D E T epsilon cubed beta squared >
Tom: Okay, so if you’re deploying this on a real model, what does that tell you about the hardware side? Does it mean we need massive amounts of extra memory or something huge for those state variables they mentioned? >
Meng: Actually, no. It says SAZO maintains MeZO’s memory footprint and only adds O(N) additional scalar state compared to MeZO, so it keeps the per-update latency low enough to actually run production systems >
Jane: That’s what I like. So, even though it's a complex adaptive strategy, the deployment cost doesn't skyrocket just because of learning sample importance dynamically >
Lu: And they also show that this method is robust to the smoothing coefficient beta, specifically performing best when beta is zero point zero five, which suggests there’s a sweet spot in how fast those importance averages update >
Tom: So, to put it simply, the naive question a person might ask is: why bother with all this complexity if standard uniform sampling works okay for some tasks? >
Jane: The answer is that when you have heterogeneous samples—when some parts of the training data are way more important than others—uniform sampling wastes evaluations on the unimportant stuff >
Lu: SAZO lets you exploit that heterogeneity by estimating per-example gradient norms from existing directional differences, which is how it figures out where to go next >
Tom: So, for someone just listening who doesn't read the paper, what does this mean for their work? Does it mean they can skip a lot of query steps and get to the target accuracy faster? >
Jane: It means they could potentially see up to six times faster convergence to the same training loss when comparing SAZO against methods like MeZO across models like Llama3-8B > <ref:2610.12566#pg2>
Lu: That speedup is what really stands out. They show this gain holds true across different model sizes and training population sizes, which suggests the underlying mechanism is quite general >
Meng: I'd be skeptical of those six times faster numbers if they only ran it on a small toy model, but since they tested four families, I'd take that as evidence that it’s not just luck >
Tom: So the paper’s conclusion is that SAZO is a sample-adaptive zeroth-order optimizer that reuses existing directional differences to estimate sample importance without needing true gradients >
Jane: And the implication for us is that we can start treating our query budget more intelligently, moving away from just random sampling towards a strategy guided by what the model actually needs at any given moment >
Lu: It’s a solid piece of theory because it connects adaptive sampling to the minimum-variance oracle, which is a very strong theoretical anchor >
Tom: Next time we talk about optimization, we'll look at how this compares against other methods like those focusing on policy updates in reinforcement learning.
The paper's improvements: Tom: We’re looking at how they suggest improving SAZO, which is all about making the system even smarter about which queries to ask next >
Jane: So, if we take what we know—that it tracks importance via moving averages—what’s the next step in making it better? >
Lu: They focus on the sampling strategy itself. They propose using this estimated importance to assign different query weights across examples to reduce that update variance even further >
Tom: That means they aren't just deciding *if* we should sample an example, but *how much* weight to give it in the actual update step >
Jane: Exactly. They formulate a new importance-weighted update equation, like gbt = sIt / N pIt, which directly uses those estimated norms to steer the optimization more precisely >
Lu: The paper shows this weighted update is ideal when the allocation is proportional to the current per-example gradient norms, which they estimate using those scaled magnitudes of existing directional differences >
Tom: So, this is just a refinement of their core idea—it’s not adding a whole new mechanism, it's just optimizing how they use the information they already have >
Jane: It sounds like this is what moves it from just "learning sample importance" to actually being an effective "sample-adaptive optimizer" that makes concrete gains >
Lu: The authors prove that this improved allocation leads to a better performance bound, asymptotically matching the leading term of the minimum-variance sampling oracle, which is a very high bar in optimization >
Tom: So they are showing that by doing this specific weighting based on those estimated norms, you get closer to the theoretical best performance without needing true gradients >
Jane: And they confirm it’s robust. They showed that even with some noise or drift in the importance estimation, like when beta is set to zero point zero five, the gains are still solid >
Lu: The analysis on robustness is important because it gives us bounds on how much error we can expect from this adaptive system across different model sizes and population sizes >
Tom: For someone who only listens to the show, what does this mean for their daily work deploying models? Do they need to rewrite their entire optimization pipeline just to use SAZO? >
Jane: It means they could swap out a simpler uniform sampler for SAZO and expect better results, potentially cutting down the total number of query steps needed by up to six times >
Lu: The implication is that heterogeneity in your training data—which is always the case with large datasets—actually becomes an asset rather than a problem when using this framework >
Tom: So, if I'm deploying this, I'm looking at maintaining that O(N) state and relying on those directional differences to give me smarter updates >
Jane: Right. And we’ll keep checking how it performs against the next set of papers we look at for comparison.
Conclusion: Tom: So we’re wrapping up this deep dive on Sample-Adaptive Zeroth-Order Post-Training, SAZO, which is about using existing directional differences to make adaptive queries online >
Jane: It really boils down to an AI system that learns sample importance dynamically so it can stop wasting queries on unimportant data >
Lu: The main thing we saw is how this method achieves near-optimal performance by matching the minimum-variance sampling oracle's leading convergence term asymptotically >
Tom: And the practical implication for me, as someone who trains models, is that if I have heterogeneous training data, I can potentially cut my query count by a factor of six compared to just using uniform sampling >
Jane: That speedup is huge. It shows that exploiting how important different parts of the training set are can lead to much better accuracy for the same computational budget >
Meng: From an engineering standpoint, it’s good because SAZO keeps MeZO’s memory footprint, so we don't have to overhaul our whole infrastructure just for a better optimization algorithm >
Lu: Exactly. The robustness analysis is solid; it shows the performance gains hold up even when there's noise or drift in the importance estimation >
Tom: I agree, the fact that it’s robust across different model families like Llama3-8B and Phi-two means this isn't just a niche trick for one specific architecture > <ref:2610.12566#pg2>
Jane: So, to summarize, SAZO is a sample-adaptive zeroth-order optimizer that reuses existing directional differences to estimate sample importance without needing true gradients >
Lu: It connects those existing differences to the ideal gradient norms using per-example exponential moving averages that track evolving importance over time >
Tom: The future work I see here is how they can integrate this adaptive sampling more tightly with the actual model fine-tuning steps, rather than just treating it as a separate optimizer component >
Jane: That sounds like a natural next step. It moves it from an external query strategy to something deeply integrated into the training loop >
Meng: For me, the cost of running this is manageable because it doesn't require extra model evaluations or huge memory overhead, which is what we worry about in production >
Lu: The paper does flag that its primary limitation is that it still relies on estimating gradient norms from directional differences rather than true gradients directly >
Tom: So, a final thought on Sample-Adaptive Zeroth-Order Post-Training: it’s a solid approach for making existing zeroth-order methods much more efficient when your data isn't uniformly useful >
Jane: It’s definitely worth keeping an eye on as we look at how we scale these optimization techniques for even larger models >
More episodes
- 2508.10173-Benchmark-Driven Selection of AI: Evidence from DeepSeek-R1
- 2602.15472-Fluids You Can Trust: Property-Preserving Operator Learning for Incompressible Flows
- 2604.09095-Geometric Probing for Algorithm Selection in Continuous Black-Box Optimization
- 2606.06459-Detecting Control and Response Events for AI-Enabled Radio Access Networks
- 2607.09816-RUBRIC: Realism--Utility Balanced Ranking for Imbalanced Classification
- 2606.04033-Gradient-based optimization of nuclear criticality experiments using neural surrogate eigenvalue sensitivities
- 2609.40265-OpenTSLM TeeMoE: A Unified Time-Series Language Model for Forecasting, Contextual Prediction, and Reasoning
- 2610.00678-Nonparametric Distribution Matching for Self-Supervised Whole-Slide Image Condensation
- 2610.02488-Harnessing LLMs as Agents: What Does It Cost?
- 2610.07061-ImpactMat: Continuous Material Estimation for Inverse Impact Sound Rendering