1 Introduction↩︎

Since their debut, Adam and related methods [1], [2] have dominated deep learning optimization. Yet, the field is now at an inflection point. Recent advances highlight a new generation of algorithms designed to better capture the geometry of modern models, with Muon [3] and its successors—Scion [4] and Gluon [5]—emerging as promising alternatives. Fueled by state-of-the-art performance in large language models (LLMs) training [6]–[10] and emerging theoretical developments [4], [5], [11], these methods are on track to disrupt entrenched practices.

Central to their design are the layer-specific linear minimization oracles (LMOs) over non-Euclidean norm balls, enabling better alignment with the highly anisotropic loss landscapes of neural networks. Concretely, let \(X = [X_1, \ldots, X_b]\) denote the parameters of a \(b\)-layer model, with \(X_i\) indexing the parameters of layer \(i \in[b] :=\{1, \ldots, b\}\). Each of the aforementioned optimizers can be viewed as an instance of the general update rule \[\begin{align} \label{eq:muon} X_i^{k+1} = X_i^k + {\rm LMO}_{\mathcal{B}_i(0,t_i^k)}(M_i^k), \qquad {\color{myred} i\in\{1, \ldots, b\}} \end{align}\tag{1}\] (see 7.1). Here, \(\mathcal{B}_i(X_i,t_i) :=\{Z_i \in \mathcal{X}_i: \left\| X_i - Z_i \right\|_{(i)} \leq t_i\}\) is a ball of radius \(t_i\) centered at \(X_i\) in the vector space \(\mathcal{X}_i\), where \(\left\| \cdot \right\|_{(i)}\) is a norm chosen for the \(i\)th layer, \(M_i^k\) is a momentum term, and the linear minimization oracle is defined as \({\rm LMO}_{\mathcal{B}_i(X_i,t_i)}(M_i) :=\mathop{\mathrm{arg\,min}}_{Z_i \in \mathcal{B}_i(X_i,t_i)} \left\langle M_i,Z_i\right\rangle\). Different choices of \(\left\| \cdot \right\|_{(i)}\) produce different algorithms (for example, Muon uses the spectral norm for hidden layers). Crucially, however, all these updates share a common characteristic: all layers are updated at every iteration. In this work, we question this design choice. Our central hypothesis is simple yet fundamental:

In the remainder of this paper, we demonstrate that the default practice of full-network updates is indeed not universally the best choice, from both theoretical and practical points of view, calling into question a core principle of standard training protocols.

1.1 Background↩︎

Let us begin by formalizing the setup. We consider the optimization problem \[\begin{align} \label{eq:problem} \min_{X\in\mathcal{X}} \left\{ f(X) :={\mathbb{E}}_{\xi\sim\mathcal{P}}\left[f(X; \xi)\right] \right\}, \end{align}\tag{2}\] where \(X \in \mathcal{X}\) represents the collection of trainable parameters of a neural network. Specifically, \(X\) is composed of block variables \(X_i \in \mathcal{X}_i :=\mathbb{R}^{m_i \times n_i}\) corresponding to layer \(i \in[b]\); we write \(X = [X_1, \ldots, X_b]\). In this context, \(\mathcal{X}\) is the \(d\)-dimensional product space \[\begin{align} \mathcal{X}:=\bigotimes_{i = 1}^b \mathcal{X}_i \equiv \mathcal{X}_1 \otimes \ldots \otimes \mathcal{X}_b, \end{align}\] where \(d :=\sum_{i=1}^b m_i n_i\). Each function \(f(\cdot; \xi): \mathcal{X}\to \mathbb{R}\) is continuously differentiable, potentially nonconvex and non-smooth, and represents the loss of the model evaluated at a data point \(\xi\) sampled from the probability distribution \(\mathcal{P}\). We denote by \(\nabla_i f(X) \in \mathcal{X}_i\) the gradient component corresponding to the \(i\)th layer, so that \(\nabla f(X) = [\nabla_1 f(X), \dots, \nabla_b f(X)] \in \mathcal{X}\). Each space \(\mathcal{X}_i\) is equipped with the trace inner product, defined as \(\langle X_i, Y_i \rangle_{(i)} :=\mathop{\mathrm{tr}}(X_i^{\top} Y_i)\) for \(X_i,Y_i \in \mathcal{X}_i\), which induces the standard Euclidean norm, denoted by \(\left\| \cdot \right\|_2\). In addition, each space is endowed with an arbitrary norm \(\left\| \cdot \right\|_{(i)}\) (which need not be induced by this inner product). We let \(\left\| \cdot \right\|_{(i) \star}\) be the dual norm associated with \(\left\| \cdot \right\|_{(i)}\) (i.e., \(\|X_i\|_{(i) \star} :=\sup_{\left\| Z_i \right\|_{(i)} \leq 1} \left\langle X_i,Z_i\right\rangle_{(i)}\) for any \(X_i\in \mathcal{X}_i\)).

2 The Algorithm↩︎

Before summarizing our main contributions (see the end of 3), we dive directly into presenting our method. Motivated by the considerations in 1, we propose Drop-Muon ([alg:rt95arbitrary95stoch])–a non-Euclidean layer-wise optimizer for deep learning based on the idea of sub-network training. At each iteration \(k\), instead of updating the entire network as in standard Muon, Drop-Muon samples a random subset \(S^k \subseteq [b]\) of layers according to a user-defined distribution \(\mathcal{D}\) and updates only the parameters of layers in \(S^k\), keeping all other layers frozen. As the reader may have noticed, the main LMO update step 3 admits two equivalent formulations. The alternative representation of 1 uses the sharp operator [12], [13], defined for any \(M \in \mathcal{X}\) by \(M^{\sharp} :=\mathop{\mathrm{arg\,max}}_{X \in \mathcal{X}} \{\left\langle M,X\right\rangle - \frac{1}{2} \left\| X \right\|^2\}\). It is well known that \(M^{\sharp}\) relates to the LMO via \(M^{\sharp} = - \left\| M \right\|_{\star} {\rm LMO}_{\mathcal{B}(0,1)}(M)\), and hence \[\begin{align} \label{eq:lmo95sharp} X_i^{k+1} = X_i^k + t_i^k {\rm LMO}_{\mathcal{B}_i(0,1)}(M_i^k) = X_i^k - \frac{t_i^k}{\left\| M_i^k \right\|_{(i) \star}} \left( M_i^k \right)^{\sharp}, \end{align}\tag{4}\] which corresponds to a layer-wise normalized steepest descent step with stepsize \(\gamma_i^k :=\frac{t_i^k}{\left\| M_i^k \right\|_{\star}}\). When \(\left\| \cdot \right\|_{(i)}=\left\| \cdot \right\|_2\) is the standard Euclidean norm, the sharp operator reduces to the identity mapping, so that \(M^{\sharp} = M\), and the update coincides with Stochastic Gradient Descent with momentum (SGDM) [14], though here performed layer-wise. These equivalent formulations will be repeatedly invoked in the proofs of the results from Sections 4.1 and 4.2.

3 Cost Model↩︎

Our key theoretical contribution is that strategically skipping some layer updates may lead to performance gains. To isolate the core phenomenon, we first study a simplified variant of the method without stochasticity in the gradients and momentum, described in 1. This deterministic version follows the same fundamental principles as [alg:rt95arbitrary95stoch], with the difference that momentum terms \(M_i^k\) are replaced by the components of the full gradient \(\nabla_i f(X^k)\). When \(\left\| \cdot \right\|_{(i)}=\left\| \cdot \right\|_2\) is the Euclidean norm, the update [eq:sharp95upd] coincides with that of layer-wise Gradient Descent (GD).

How expensive is one step of 1? The answer is governed by the sampling distribution \(\mathcal{D}\). Consider iteration \(k\) and denote \(s^k :=\min S^k\) (the smallest index of an active layer). The operations performed by the algorithm can be summarized as follows:

1. Backward pass: Backpropagate gradients through layers \([X^k_{s^k}, \ldots, X^k_b]\). Since layers \(1, \ldots, s^k - 1\) are frozen, no gradients are computed for them, effectively truncating the gradient flow at the first active layer.

2. Forward pass: To evaluate the loss, activations must in principle be propagated through all layers \(1, \ldots, b\). However, since only layers \([X^k_{s^k}, \ldots, X^k_b]\) are updated at iteration \(k\), the activations up to layer \(s^k - 1\) may be cached and reused in the next step.

3. Gradient transformation: Given the gradients \(\{\nabla_i f(X^k)\}_{i\in S^k}\), compute the corresponding sharp operators \(\{(\nabla_i f(X^k))^{\sharp}\}_{i\in S^k}\) (or, equivalently, the LMOs; see 4 ).

4. Parameter updates: Update the parameters of layers \(\{X_i^k\}_{i\in S^k}\) using their computed (transformed) gradients, while keeping the frozen layers unchanged.

To model the total computational effort of the optimization procedure, we associate a cost with each step (measured, for example, in FLOPs or wall-clock time). Let \(c_{\mathrm{ov}} \geq 0\) denote the fixed per-iteration overhead (e.g., data loading). As noted above, backpropagation must be performed from the last layer \(b\) down to layer \(s^k\), while forward-pass activations up to layer \(s^k - 1\) can be cached and reused in subsequent iterations. Hence, the costs of steps (i) and (ii) can be aggregated into a single per-layer constant \(c_i > 0\) for each \(i \in [b]\). In the non-Euclidean setting (where \(M^{\sharp} \neq M\)), an additional cost arises from computing sharp operators. We denote by \(c_i^{\sharp} \geq 0\) the combined cost of evaluating this operator and performing the corresponding parameter update for layer \(i\) (steps (iii) and (iv)). Under this model, the total compute cost of iteration \(k\) is \[\begin{align} \label{eq:cost95k} \mathrm{cost}(S^k) :=c_{\mathrm{ov}} + \sum\limits_{i=s^k}^b c_i + \sum\limits_{i \in S^k} c_i^{\sharp}, \end{align}\tag{5}\] and, consequently, for a fixed target accuracy \(\varepsilon > 0\), the expected cost of the entire optimization procedure can be expressed as \[\begin{align} \label{eq:exp95cost95gen} \mathrm{cost}_{\varepsilon}(\mathcal{D}) :=K \times {\mathbb{E}}\left[\mathrm{cost}(\hat{S})\right], \end{align}\tag{6}\] where we write \(\hat{S} \sim \mathcal{D}\) to denote a random variable with the same distribution as that of the samplings (since \(\{S^k\}_{k\geq0}\) are i.i.d.), and \(K\) is the number of iterations to reach convergence (interpreted in the nonconvex case as reaching an \(\varepsilon\)-stationary point in expectation). In the remainder of this paper, we compute \(K\) under two smoothness regimes: layer-wise smoothness (2) and layer-wise \((L^0, L^1)\)–smoothness (3). We then evaluate various layer-update strategies and compute their expected costs as in 6 . These results converge to a single conclusion:

We provide a rigorous statement and justification for this claim in the sections that follow. The main takeaway, however, is simple: as this condition is highly unlikely to be realized in practice, updating only a subset of layers at each iteration should be seen as more efficient than the default strategy of updating all parameters at each iteration.

Our contributions can be summarized as follows:

1. Challenging full network updates. We provide, to our knowledge, the first systematic investigation of the practice of updating all layers of a network at every iteration. We argue—and rigorously demonstrate both in theory and in practice—that this design choice is generally suboptimal.

2. General framework for sub-network optimization. We introduce Drop-Muon ([alg:rt95arbitrary95stoch] and its deterministic gradient counterpart 1), a principled layer-wise optimization framework with randomized layer subsampling. Drop-Muon strictly generalizes LMO-type methods (including Muon [3], Scion [4], and Gluon [5]) by allowing random subsets of layers to be updated per step, with full-network training as a special case. Drop-Muon supports virtually any layer sampling scheme (9). In the main part of this paper, we focus on Randomized Progressive Training (RPT) [15], a natural strategy aligned with backpropagation mechanics that avoids redundant gradient computations and reduces compute cost while maintaining strong convergence guarantees.

3. Tight iteration complexity guarantees under novel smoothness regimes. We establish convergence guarantees for Drop-Muon under two regimes: layer-wise smoothness (1) and layer-wise \((L^0,L^1)\)–smoothness (Theorems 2 and 4). Our rates recover the state of the art for SGD- and Muon-type methods, and, to our knowledge, provide the first convergence guarantees for progressive training-style methods in the non-smooth setting.

4. Theoretical compute-optimality results. To isolate the key phenomena, we first consider a deterministic gradient variant of Drop-Muon (1). Using a simple yet expressive cost model (5 ) accounting for per-layer forward/backward passes, gradient transformations, and parameter updates, we prove that full-network updates are not optimal unless a very specific condition on layer smoothness constants holds (Theorems 3 and 15), which is unlikely in practice. This formally justifies selective layer updates as the compute-optimal default.

5. Empirical validation. Controlled CNN experiments on MNIST, Fashion-MNIST, and CIFAR-10 show that Drop-Muon consistently outperforms standard full-network Muon, achieving the same accuracy up to \(1.4\times\) faster in wall-clock time.

4 Randomized Progressive Training↩︎

The general framework in 1 allows virtually any sampling strategy. However, due to the mechanics of backpropagation, it is most natural to update all layers from the last one down to some sampled minimal index. Specifically, if the smallest sampled index at iteration \(k\) is \(s^k\), then computing the gradient \(\nabla_{s^k} f(X^k)\) requires backpropagating from the last layer \(b\) up to layer \(s^k\), which automatically produces all gradient components \([\nabla_{s^k} f(X^k), \ldots, \nabla_b f(X^k)]\).

Formally, we can define the sampling distribution \(\mathcal{D}\) as follows: at each iteration \(k\), sample \(s^k \in [b]\) with probabilities \(p_i :={\mathbb{P}}\left(s^k=i\right)\), where \(\sum_{i=1}^b p_i = 1\) and \(p_1>0\), and set \(S^k = \{s^k, \ldots, b\}\). Algorithms [alg:rt95arbitrary95stoch] and 1 then update the layers \([X^k_{s^k}, \ldots, X^k_b]\), while \([X_1^k, \ldots, X_{s^k-1}^k]\) remain frozen at their previous values (with the convention that \([X_1^k, X_0^k] = \emptyset\) when \(s^k = 1\)). We refer to this sampling scheme as Randomized Progressive Training (RPT), or Drop-training for short (see 5).

4.1 Iteration Complexity – Deterministic Gradient Setting↩︎

We now analyze the iteration complexity of 1 under two smoothness regimes; the proofs are deferred to the Appendix. Throughout, we make the standard assumption that the objective function is lower-bounded.

This ensures the existence of an approximately stationary point for any desired level of accuracy.

4.1.0.1 Smooth case.

We first establish convergence under the layer-wise smoothness assumption.

We take \(L^0\) to be the smallest collection of constants satisfying the above. Throughout, we use \(\textrm{supp}(\mathcal{D})\) to denote the subsets of \([b]\) assigned positive probability mass by \(\mathcal{D}\). In the progressive training setting, these supported sets take the form \(\textrm{supp}(\mathcal{D}) = \{\{j,\ldots,b\}, j\in[b]\}\). By definition, \(L_{i,S}^0 = 0\) whenever \(i \notin S\). Moreover, if \(S_1, S_2 \in \textrm{supp}(\mathcal{D})\) are such that \(S_1 \subseteq S_2\), then \(L_{i,S_1}^0 \le L_{i,S_2}^0\) (see 1).

The assumption is inspired by the coordinate descent (CD) literature [16], reducing to the standard block-wise Lipschitz continuity of the gradient [17]–[19] in the special case when \(\textrm{supp}(\mathcal{D}) = \{\{1\}, \{2\}, \ldots, \{b\}\}\). 2 captures the intuition that each subset of layers can have its own effective smoothness constant, allowing tighter bounds on the local curvature of \(f\) and better reflecting the structure of the model. Importantly, it is not more restrictive than standard smoothness–rather, it offers a richer parametric description by assigning separate constants to different layer subsets, allowing a more precise analysis without shrinking the function class.

With the assumptions in place, we are now ready to state the first formal convergence result.

1 establishes an \(\mathcal{O}(K^{-1})\) convergence rate for a weighted sum of squared gradient component norms, matching the theoretical rates previously established for Muon-type methods under classical smoothness [4], [11], [20]. For the result to be meaningful, every gradient component must contribute to the weighted average. In other words, the sampling distribution must satisfy \(w_i>0\) for all \(i\in[b]\). This is a natural requirement, equivalent to ensuring that \(p_1>0\), i.e., that all layers are updated with nonzero probability. Obviously, if this was not the case, the first layer would be completely ignored, making convergence impossible.

4.1.0.2 Generalized smooth case.

Layer-wise optimizers considered here are designed for deep learning, where the classical smoothness assumption is often violated [5], [21]. Consequently, the layer-wise smoothness model in 2 may not accurately capture the local geometry of the loss. To address this, we adopt a more expressive framework building upon \((L^0, L^1)\)–smoothness [21], [22]. 3 below generalizes 2 by letting the local curvature of each layer depend not only on fixed constants \(L_{i,S}^0\), but also on the magnitude of the layer’s gradient via additional terms \(L_{i,S}^1 \left\| \nabla_i f(X) \right\|_{(i) \star}\).

As with 2, assigning separate constants to each subset of layers \(S\) allows for tighter, subset-specific bounds, reflecting the interactions among layers.

Similar to 1, 2 guarantees an \(\mathcal{O}(K^{-1/2})\) convergence rate for a weighted sum of gradient component norms.² Again, the sampling distribution must ensure that \(w_i > 0\) for all \(i\in[b]\), which amounts to requiring that \(p_1>0\). In the extreme case when \((p_1,p_2,\ldots,p_b) = (1,0,\ldots,0)\), corresponding to full-network training, the weights simplify to \(w_i = \frac{1}{L^1_{i,[b]}}\), exactly recovering the convergence rate of deterministic Gluon [5], demonstrating the tightness of our guarantees. Importantly, the stepsizes naturally scale inversely with the layer-specific smoothness constants and gradient magnitudes. This automatic adaptation to local geometry prevents overshooting, ensures stable convergence, and allows more aggressive updates when the \((L^0, L^1)\) constants and gradient norms are small.

4.1.1 Cost Optimization↩︎

Let us now make clear why performing full-network updates in all iterations is not in general an optimal strategy. We first consider the layer-wise smooth case. According to 1, under 2, 1 guarantees that \(\frac{1}{K} \sum_{k=0}^{K-1} \sum_{i=1}^b {\mathbb{E}}\left[\left\| \nabla_i f(X^k) \right\|^2_{(i) \star}\right] \leq \varepsilon\) after \[\begin{align} K = \left\lceil \frac{f(X^0) - f^{\star}}{\varepsilon \times \min_{i\in[b]} \left[\sum\limits_{s=1}^i \frac{p_s}{2 L_{i, \{s,\dots,b\}}^0}\right]} \right\rceil \end{align}\] iterations. Since in this case \({\mathbb{E}}\left[\mathrm{cost}(\hat{S})\right] = c_{\mathrm{ov}} + \sum_{j=1}^b (c_j + c_j^{\sharp}) \sum_{i=1}^j p_i\), the expected total cost of the optimization procedure satisfies \[\begin{align} \label{eq:smooth95cost95full} \mathrm{cost}_{\varepsilon}(\mathcal{D}) \,\overset{\eqref{eq:exp95cost95gen}}{=}\, K \times {\mathbb{E}}\left[\mathrm{cost}(\hat{S})\right] \,\propto\, \frac{c_{\mathrm{ov}} + \sum_{i=1}^b (c_i + c_i^{\sharp}) \sum_{s=1}^i p_s}{\min_{i\in[b]} \left[\sum\limits_{s=1}^i \frac{p_s}{2 L_{i, \{s,\dots,b\}}^0}\right]} \end{align}\tag{7}\] (see 11.1.1 for the details of the derivation). Finding the sampling distribution that minimizes \(\mathrm{cost}_{\varepsilon}(\mathcal{D})\) is equivalent to optimizing over \(\{p_i\}_{i\in[b]}\). Letting \((p_1,p_2,\ldots,p_b) = (1,0,\ldots,0)\) recovers full-network training, which serves as a natural baseline. Yet, can we do better? The following theorem shows that this configuration is optimal under very specific conditions only.

Note that the condition in 3 is entirely independent of the cost parameters! In fact, one can derive a recursive construction of the optimal probabilities that depends solely on the smoothness constants, from which 3 follows as a simple corollary (see 14).

The layer-wise \((L^0, L^1)\)–smooth case (15) leads to the same conclusion: full-network updates are optimal only if the first layer is associated with the largest smoothness constant–a restrictive and rarely observed condition, as confirmed by experiments on NanoGPT [5]. While broader validation is needed, there is no reason to expect this phenomenon to hold broadly. Overall, from the theoretical standpoint, the prevalent practice of always updating all layers is fundamentally flawed.

4.2 Iteration Complexity – Stochastic Gradient Setting↩︎

We now turn to the convergence analysis of the practical variant of Drop-Muon with stochastic gradients and momentum ([alg:rt95arbitrary95stoch]), within the general layer-wise \((L^0, L^1)\)–smoothness framework.

4 is a slightly stronger variant of 3 used in the analysis of 1 (see 10). Consequently, all results proven under 3 carry over to this setting. This stronger form is necessary to rigorously establish convergence in the stochastic case.

Furthermore, we assume access to a standard stochastic gradient oracle with bounded variance.

Note that, to facilitate the proofs, the variance bound in 5 is defined with respect to the Euclidean norm, consistent with prior analyses [4], [5], [11]. This is not restrictive: since \(\mathcal{X}\) is finite-dimensional, there exist \(\underline{\rho}_i, \bar{\rho}_i > 0\) such that \(\underline{\rho}_i \left\| X_i \right\|_{(i)} \leq \left\| X_i \right\|_2 \leq \bar{\rho}_i \left\| X_i \right\|_{(i)}\) for all \(X_i\in\mathcal{X}_i\) (equivalently, \(\underline{\rho}_i \left\| X_i \right\|_2 \leq \left\| X_i \right\|_{(i) \star} \leq \bar{\rho}_i \left\| X_i \right\|_2\)).

4 establishes an \(\mathcal{O}(K^{-1/4})\) convergence rate for a weighted sum of gradient component norms, in line with the state-of-the-art results for SGD- and Muon-type methods [5], [11], [14], [23]. As in the deterministic setting, one could attempt a cost analysis; however, the stochastic bound’s complexity necessitates a case-by-case treatment depending on which term dominates. While this prevents deriving as clean analytic results as in 4.2, our experiments (6) clearly confirm that the approach remains effective in the stochastic setting, too.

5 Prior Work on Progressive Training↩︎

The concept of partial network updates has been explored in several prior works. Progressive Training (PT) [24] was first introduced in the context of Generative Adversarial Networks (GANs) [25]. The central idea is to begin with a shallow network trained on low-resolution inputs and gradually increase both the input resolution and network depth by adding layers over time. This progressive growing strategy has been shown to improve computational efficiency and stabilize the training process. Despite its empirical success, PT has remained a largely heuristic method without formal convergence guarantees. An early attempt to provide such guarantees was made by [26] with ProgFed, which also extended progressive training to distributed settings. However, the accompanying theoretical analysis has been criticized as vacuous [15]. A more rigorous treatment was proposed by [15], who introduced Randomized Progressive Training (RPT). This method can be viewed as a randomized proxy for PT and was the first to establish theoretical convergence rates for progressive training on general smooth objectives, building on the framework of Randomized Coordinate Descent (RCD) [12], [27]. Conceptually, RPT is closely related to our approach: like Drop-Muon, it updates a random subnetwork at each iteration. However, the method of [15] is simply a form of sketched GD with a particular choice of sketch operator, and has several limitations that restrict its practical utility. First, it treats network parameters as a flat vector, ignoring the layer-wise structure. Consequently, it fails to exploit layer-specific, non-Euclidean update rules–failing, to recover methods such as Muon as special cases. Second, RPT requires non-stochastic gradient computation, making it computationally infeasible for large-scale architectures.

6 Experiments↩︎

We evaluate Drop-Muon on three benchmarks–MNIST, Fashion-MNIST, and CIFAR-10–using 3-layer convolutional neural networks (CNNs) of varying capacity. For all \(i \in [b]\), we set \(\left\| \cdot \right\|_{(i)} = \left\| \cdot \right\|_{2\to2}\), i.e., the spectral norm, consistent with Muon. We study RPT sampling strategy with several smallest-index sampling rules, including uniform and an epoch-shift distribution that gradually shifts probability mass from shallow to deep layers as training progresses (see 13). Our baseline is standard Muon (i.e., Drop-Muon with spectral norms and full-network training) run under identical hyperparameter configurations. This design allows us to isolate the effect of partial network updates while controlling for all other factors. For each configuration, we record accuracy as a function of both epochs and wall-clock time. Further experiments and details of hyperparameter selection procedure are provided in 13.

6.0.0.1 Results on MNIST.

2 illustrates a representative run with uniform layer index sampling: although Drop-Muon converges slightly slower per epoch (left), its lower per-epoch cost (\(27\)s vs. \(39\)s for Muon) yields faster overall training in wall-clock time (right). This advantage is even clearer with epoch-shift sampling: 3 (left) aggregates time-to-target results over multiple seeds, showing that Drop-Muon consistently outperforms Muon and achieves up to \(1.4\times\) speed-up.

6.0.0.2 Results on Fashion-MNIST.

We repeat the experiment on Fashion-MNIST. Drop-Muon again delivers meaningful acceleration: as before, its per-epoch convergence is slightly slower, but it overtakes Muon in wall-clock time. 3 (right) shows aggregated time-to-target results over multiple seeds, with Drop-Muon achieving roughly \(1.2\times\) faster convergence on average across accuracy thresholds (we omit the \(99\%\) threshold since neither method reaches it).

6.0.0.3 Results on CIFAR-10.

4 presents a representative run. Although the absolute gain is smaller than on MNIST or Fashion-MNIST, Drop-Muon still reaches \(90\%\) training accuracy earlier in wall-clock time (see also 11).

In summary, even though Drop-Muon trains on partial gradients, it reaches high accuracy levels earlier than Muon in wall-clock time. Importantly, Drop-Muon is simple to implement, requiring only a few lines of code. The observed speed-up over Muon could be even greater with dedicated tuning: in our experiments, both methods use the same constant learning rates, whereas theory (Theorems 1, 2, and 4) and the practice of coordinate descent (7.3) suggest that Drop-Muon can safely employ larger stepsizes than the full-network Muon baseline. Here, we kept the learning rates identical for both optimizers to isolate the wall-clock benefit of partial layer updates, but additional gains are likely achievable with method-specific tuning. Overall, our results demonstrate that Drop-Muon is an effective, practical, and easily implementable strategy that, with a well-chosen sampling strategy, consistently accelerates training while retaining high accuracy, making it a compelling choice for modern neural network optimization.

Acknowledgments↩︎

The research reported in this publication was supported by funding from King Abdullah University of Science and Technology (KAUST): i) KAUST Baseline Research Scheme, ii) CRG Grant ORFS-CRG12-2024-6460, and iii) Center of Excellence for Generative AI, under award number 5940.

Appendix↩︎

7 Additional Literature Review↩︎

7.1 Muon (and Friends)↩︎

Recent work has revisited how neural network parameters are updated, moving beyond simple element-wise gradient steps toward more structured update rules. Notable examples include the Muon optimizer [3], as well as the closely related Scion [4] and Gluon [5]. All of these methods can be interpreted as instances of algorithms driven by a linear minimization oracle (LMO) over norm balls.

Muon, introduced by [3], is an optimizer for the hidden layers of neural networks (with first and last layers typically trained using AdamW [2]). Given a layer \(X_i\) and the associated (stochastic) gradient \(G_i\), the update is obtained by solving the constrained optimization problem \[\begin{align} \label{eq:muon95lmo} \min_{\Delta X_i}\;\left\langle G_i,\Delta X_i\right\rangle \quad \text{subject to} \quad \|\Delta X_i\|_{2\to 2} \le t_i, \end{align}\tag{8}\] where \(t_i > 0\) plays the role of a trust-region radius or stepsize. The solution is characterized by the singular vectors of \(G_i\): if \(G_i = U_i \Sigma_i V_i^\top\) is its singular value decomposition, then \[\begin{align} \label{eq:muon95sol} \Delta X_i = -t_i U_i V_i^\top, \qquad X_i^{k+1} = X_i^k + \Delta X_i. \end{align}\tag{9}\] In other words, Muon moves in the direction of steepest descent measured in the spectral norm.

Since computing an exact SVD can be expensive, Muon uses the Newton-Schulz method [28], [29] to approximate the orthogonalization. In practice, the algorithm also incorporates momentum, yielding the update \[\begin{align} M_i^k &= (1-\beta_i) M_i^{k-1} + \beta_i G_i^k, \\ X_i^{k+1} &= X_i^k - t_i^k \mathrm{NewtonSchulz}(M_i^k), \end{align}\] where \(\beta_i \in (0,1]\) is the momentum parameter and \(M_i^k\) is the running average of past gradients.

The update rule in 9 can be interpreted as a special case of the more general LMO-based update described in 1. By selecting different norms for the LMO ball constraint, one obtains different algorithmic variants. Since our primary focus is on deep learning applications, we are especially interested in matrix norms. Within this setting, a particularly important family is that of operator norms, defined for any \(A \in \mathbb{R}^{m \times n}\) as \[\left\| A \right\|_{\alpha \to \beta} :=\sup _{\left\| Z \right\|_\alpha=1} \left\| A Z \right\|_\beta,\] where \(\left\| \cdot \right\|_{\alpha}\) and \(\left\| \cdot \right\|_{\beta}\) denote norms on \(\mathbb{R}^{n}\) and \(\mathbb{R}^{m}\), respectively.

In particular, Muon’s update in 8 corresponds to an LMO over the spectral norm ball \(\mathcal{B}_i^{2\to2}(0,t_i) :=\{Z_i \in \mathcal{X}_i : \left\| Z_i \right\|_{2\to 2} \leq t_i\}\), leading to the equivalent update \[\begin{align} X_i^{k+1} = X_i^k + {\rm LMO}_{\mathcal{B}_i^{2\to 2}(0,t_i)}(G_i). \end{align}\] Hence, Muon is just one member of a broader family of methods parameterized by the geometry of the update set.

Building on this perspective, [4] introduced Scion, which extends LMO-based updates to all network layers, rather than being limited to the hidden matrix-shaped layers as in Muon. To better capture layer-specific behavior, Scion employs different norms depending on layer type (scaled spectral norms for the weight matrices of transformer blocks and the \(\left\| \cdot \right\|_{1\to \infty}\) norm for embedding and output layers).

Finally, Gluon [5] generalizes the theory behind both Muon and Scion. It provides a convergence framework for updates over arbitrary norm balls. The analysis relies on a novel layer-wise smoothness model (closely related to our Assumptions 2, 3, and 4), capturing heterogeneity through \((L^0,L^1)\) parameters, that more accurately reflects the non-uniform characteristics of deep learning models (see more details in 7.2). As such, Gluon can be seen as the theoretical foundation unifying the entire class of LMO-based optimizers for training deep networks.

7.2 Generalized Smoothness↩︎

Gradient-based methods are traditionally analyzed under the assumption of Lipschitz smoothness of the gradient. Yet, many modern deep learning tasks violate this assumption [21], [30], prompting the development of alternative smoothness models. One such model is \((L^0, L^1)\)–smoothness, initially introduced by [21] for twice continuously differentiable functions and later generalized beyond this setting [22], [31].

This framework has been further tailored to deep learning through the non-Euclidean layer-wise \((L^0, L^1)\)–smoothness assumption [5], which closely relates to our Assumptions 2, 3, and 4. Concretely, [5] assume that \[\begin{align} \label{eq:oagsbv} \left\| \nabla _i f(X) - \nabla _i f(Y) \right\|_{(i) \star} \leq \left( L^0_i + L^1_i \| \nabla _i f(X) \|_{(i) \star} \right) \left\| X_i - Y_i \right\|_{(i)} \end{align}\tag{10}\] for all \(i\in[b]\) and all \(X = [X_1, \ldots, X_p]\in \mathcal{X}\), \(Y = [Y_1, \ldots, Y_p] \in \mathcal{X}\). Condition 10 can be interpreted as a global version of our 4. However, it is less precise because it does not allow each subset of layers to have its own effective smoothness constants. By adopting 4 instead of 10 , we can derive tighter bounds without restricting the function class under consideration.

Accounting for heterogeneous parameter structures is not novel and has been studied in the context of coordinate descent [12], [27] and in the analysis of algorithms such as signSGD [30], [32], AdaGrad [33], [34], and Adam [35].

7.3 Randomized Block Coordinate Descent↩︎

Coordinate Descent (CD) algorithms have a long history [12], [16], [36], [37] and are among the most widely studied methods for large-scale optimization. Their randomized variants, known as Randomized Block Coordinate Descent (RBCD), have emerged as a powerful class of first-order methods, particularly in high-dimensional problems. Instead of updating the entire parameter vector, RBCD updates only a subset (“block”) of variables at each iteration, reducing per-iteration computational cost and improving scalability.

In the context of CD, the parameters are viewed as a flat vector in \(\mathbb{R}^d\). The general procedure of RBCD is summarized in 5. At each iteration, one block is selected at random and updated, while the remaining blocks remain unchanged. Clearly, 5 is a special case of 1 with the serial sampling strategy (i.e., sampling a single coordinate per iteration) and standard Euclidean updates.

We define a family of linear maps \(\{\mathbf{U}_i: \mathcal{X}_i\rightarrow \mathcal{X}: i\in [b]\}\) (where \(\mathcal{X}=\mathbb{R}^d\) and \(\mathcal{X}_i=\mathbb{R}^{d_i}\), with \(\sum_{i=1}^b d_i = d\)), forming a partition of the identity map, i.e., \(X = \sum_{i\in [b]} \mathbf{U}_i \left(X_i\right)\). The standard analysis of RBCD relies on the assumption of blockwise Lipschitz continuity of the gradient: \[\label{eq:sdfewr} \left\| \nabla_i f(X+\mathbf{U}_i (\Gamma_i)) - \nabla_i f(X) \right\|_{(i)\star}\leq L_i \|\Gamma_i\|_{(i)},\quad \forall i\in [b], \, X \in \mathcal{X}, \, \Gamma_i \in \mathcal{X}_i.\tag{11}\] From 11 , one obtains the standard quadratic upper bound \[\label{eq:fnf} f(X+ \mathbf{U}_i (\Gamma_i))-f(X)- \left\langle \nabla_i f(X), \Gamma_i \right\rangle \leq \frac{L_i}{2} \|\Gamma_i\|^2_{(i)}.\tag{12}\] This inequality plays a pivotal role in establishing the iteration complexity of RBCD. Moreover, it explains why CD-type methods can be advantageous compared to full gradient descent: they allow for larger stepsizes, which in turn leads to faster convergence. Indeed, the standard analysis of GD relies on global \(L\)-smoothness, enforcing stepsizes on the order of \(\frac{1}{L}\). By contrast, the blockwise smoothness assumption 11 bounds the function’s variation when only a single block of coordinates is updated. Consequently, in CD, the safe stepsize for block \(i\) scales with \(\frac{1}{L_i}\). Since \(L\) must capture worst-case variations across all directions (including cross-block interactions), it is often much larger than the individual \(L_i\)’s.

Complexity analyses of RBCD in the convex setting can be found in [12], [16]. Generalizations to parallel and arbitrary sampling variants were developed by [19], [38]. For the nonconvex case considered in this paper, iteration complexity results of 1 are only known for serial sampling [39], [40]. Otherwise, some work has focused on the global convergence properties of randomized CD methods [41], rather than the iteration complexity analysis we develop here. It is worth noting that serial sampling is not suitable in the deep neural network context due to the special entanglement structure and the requirements of backpropagation for gradient computation.

8 More on Smoothness Assumptions↩︎

Recall that a function \(f: \mathcal{X}\mapsto \mathbb{R}\) is \(\textrm{supp}(\mathcal{D})\)–layer-wise \(L^0\)–smooth with constants \(L^0 :=\{(L_{1,S}^0, \ldots, L_{b,S}^0)\}_{S\in\textrm{supp}(\mathcal{D})}\), \((L_{1,S}^0, \ldots, L_{b,S}^0) \in \mathbb{R}^b_+\) (2) if for any \(S\in\textrm{supp}(\mathcal{D})\), \[\begin{align} f(X + \Gamma) - f(X) - \left\langle\nabla f(X),\Gamma\right\rangle \leq \sum_{i\in S} \frac{L_{i,S}^0}{2}\left\| \Gamma_i \right\|_{(i)}^2 \end{align}\] for all \(X = [X_1, \ldots, X_b]\in \mathcal{X}\) and \(\Gamma = [\Gamma_1, \ldots, \Gamma_b] \in \mathcal{X}\) such that \(\Gamma_i = 0\) for all \(i\not\in S\).

9 Arbitrary Sampling↩︎

One might wonder why the cost model in 3 distinguishes between two sets of constants \(\{c_i\}_{i\in[b]}\) and \(\{c_i^\sharp\}_{i\in[b]}\), given that in the RPT case, \(c_i\) and \(c_i^\sharp\) could be combined into a single constant. This is because Algorithms [alg:rt95arbitrary95stoch] and 1 do not restrict us to the layer sampling scheme studied in 4. One could argue that deviating from this progressive training framework is inefficient, as it discards gradients obtained “for free” during backpropagation: by design, at iteration \(k\), the gradients \([\nabla_{s^k} f(X^k), \ldots, \nabla_b f(X^k)]\), where \(s^k :=\min S^k\), are necessarily computed, and thus available for the update (if, for some block \(i \in [b]\), the cost \(c_i^{\sharp}\) of applying the sharp operator is large, one can simply bypass this step and instead use the standard gradient update; this corresponds to setting \(\left\| \cdot \right\|_{(i)}=\left\| \cdot \right\|_2\) in the algorithm). Nevertheless, since the general iteration complexity results in Theorems 7, 10 and 16 hold for any distribution \(\mathcal{D}\), the framework naturally allows experimentation with alternative sampling schemes, even if only out of theoretical curiosity.

With this in mind, we generalize the results from 4 to arbitrary layer samplings. Formally, at each iteration \(k\), Algorithms [alg:rt95arbitrary95stoch] and 1 sample a random subset of layers \(S^k \subseteq [b]\) from a distribution \[\begin{align} \mathcal{D}: \mathfrak{P}([b]) \to [0,1], \qquad \sum_{S\subseteq[b]} \mathcal{D}(S) = 1, \end{align}\] with support \(\textrm{supp}(\mathcal{D}) :=\{S\in\mathfrak{P}([b]): \mathcal{D}(S)>0\}\), where \(\mathfrak{P}([b])\) denotes the power set of \([b]\). Only the layers in \(S^k\) are updated, while the rest remain fixed. We write \(\hat{S} \sim \mathcal{D}\) for a set-valued random variable with the same distribution as \(S^k\). This formulation defines a broad family of algorithms, parameterized jointly by the sampling distribution \(\mathcal{D}\) and the choice of layer-wise norms \(\left\| \cdot \right\|_{(i)}\).

Mirroring the earlier discussion, we organize the analysis around two smoothness regimes–layer-wise smooth (2) and generalized layer-wise smooth (Assumptions 3 and 4).

In what follows, we present convergence results for Algorithms [alg:rt95arbitrary95stoch] (16) and 1 (Theorems 7 and 10). When specialized to the RPT setting, these general results recover the guarantees stated in Theorems 1, 2, and 4, as detailed in Remarks 9, 12, and 18.

10 Convergence Results – Deterministic Gradient Case↩︎

10.1 Layer-Wise Smooth Case↩︎

10.2 Layer-Wise Generalized Smooth Case↩︎

11 Optimizing the Cost – Deterministic Gradient Case↩︎

Let \(S^k\subseteq\{1,\dots,b\}\) be the random subset sampled at iteration \(k\), and \(s^k :=\min S^k\) its smallest index. Recall the per-iteration cost \[\begin{align} \mathrm{cost}(S^k) = c_{\mathrm{ov}} + \sum_{i=s^k}^b c_i + \sum_{i\in S^k} c_i^{\sharp}. \end{align}\] Define the two marginal probabilities \[\begin{align} F_i :={\mathbb{P}}\left(\hat{s} \leq i\right),\qquad Q_i :={\mathbb{P}}\left(i \in \hat{S}\right), \end{align}\] where \(\hat{s}\) is a random variable following the same distribution as \(s^k\) (since \(S^k \sim \mathcal{D}\), \(k \geq 0\), are i.i.d., the same holds for \(s^k\)). Since \[\begin{align} {\mathbb{E}}\left[\sum_{i=s^k}^b c_i\right] = \sum_{j=1}^b {\mathbb{P}}\left(s^k = j\right) \sum_{i=j}^b c_i = \sum_{i=1}^b c_i \sum_{j=1}^i {\mathbb{P}}\left(s^k = j\right) = \sum_{i=1}^b c_i {\mathbb{P}}\left(s^k \leq i\right) \end{align}\] and \[\begin{align} {\mathbb{E}}\left[\sum_{i\in S^k} c_i^{\sharp}\right] = {\mathbb{E}}\left[\sum_{i=1}^b \mathbb{I}\left(i\in S^k\right) c_i^{\sharp}\right] = \sum_{i=1}^b c_i^{\sharp} {\mathbb{P}}\left(i\in S^k\right), \end{align}\] the expected cost of one iteration is \[\begin{align} \label{eq:cost95gen} {\mathbb{E}}\left[\mathrm{cost}(\hat{S})\right] = c_{\mathrm{ov}} + \sum_{i=1}^b c_i F_i + \sum_{i=1}^b c_i^{\sharp} Q_i. \end{align}\tag{16}\] Hence, to evaluate the expected cost, it suffices to compute the two marginals \(F_i\) and \(Q_i\) under the chosen sampling scheme.

In the sequel, we describe a few example sampling strategies considered in this work and analyze their costs under the layer-wise smooth setting (see 11.1) and the generalized layer-wise smooth setting (see 11.2).

1. RPT. Sample \(\hat{s}\in\{1,\dots,b\}\), where \(p_i = {\mathbb{P}}\left(\hat{s}=i\right)\), and set \(\hat{S}=\{\hat{s},\dots,b\}\). Then \[\begin{align} F_i = {\mathbb{P}}\left(\hat{s}\leq i\right) = \sum_{j=1}^i p_j, \qquad Q_i = {\mathbb{P}}\left(i\in \hat{S}\right) = {\mathbb{P}}\left(\hat{s}\leq i\right) = F_i. \end{align}\] Therefore, \[\begin{align} {\mathbb{E}}\left[\mathrm{cost}(\hat{S})\right] = c_{\mathrm{ov}} + \sum_{i=1}^b (c_i + c_i^{\sharp}) F_i = c_{\mathrm{ov}} + \sum_{i=1}^b (c_i + c_i^{\sharp}) \sum_{j=1}^i p_j. \end{align}\]

2. \(\tau\)-nice sampling. Choose \(\hat{S}\) uniformly from all subsets of \([b]\) of size \(\tau\). Then \[\begin{align} F_i &=& 1 - {\mathbb{P}}\left(\hat{s} > i\right) = 1 - \frac{\binom{b-i}{\tau}}{\binom{b}{\tau}}, \\ Q_i &=& \frac{\binom{b-1}{\tau-1}}{\binom{b}{\tau}} = \frac{\tau}{b}, \end{align}\] where we use the convention \(\binom{n}{k}=0\) for \(n<k\) (so the formula for \(F_i\) automatically gives \(F_i=1\) when \(b-i<\tau\)). Hence \[\begin{align} {\mathbb{E}}\left[\mathrm{cost}(\hat{S})\right] = c_{\mathrm{ov}} + \sum_{i=1}^b c_i \left( 1 - \frac{\binom{b-i}{\tau}}{\binom{b}{\tau}} \right) + \frac{\tau}{b} \sum_{i=1}^b c_i^{\sharp}. \end{align}\]

3. \(\tau\)-submodel sampling. Sample a starting index \(\hat{s}\in\{1,\dots,b-\tau+1\}\) with probability \(p_i = {\mathbb{P}}\left(\hat{s}=i\right)\) (where \(\tau\in[b]\) is fixed), and set \(\hat{S}=\{\hat{s},\dots,\hat{s}+\tau-1\}\) (i.e., a block of \(\tau\) consecutive layers). Then the marginals are \[\begin{align} F_i &=& {\mathbb{P}}\left(\hat{s} \leq i\right) = \sum_{j=1}^{\min\{i,b-\tau+1\}} p_j, \\ Q_i &=& {\mathbb{P}}\left(i\in \hat{S}\right) = \sum_{j=\max\{1,i-\tau+1\}}^{\min\{i,b-\tau+1\}} p_j. \end{align}\] The expected per-iteration cost is therefore \[\begin{align} {\mathbb{E}}\left[\mathrm{cost}(\hat{S})\right] = c_{\mathrm{ov}} + \sum_{i=1}^b c_i \left( \sum_{j=1}^{\min\{i,b-\tau+1\}} p_j \right) + \sum_{i=1}^b c_i^{\sharp} \left( \sum_{j=\max\{1,i-\tau+1\}}^{\min\{i,b-\tau+1\}} p_j \right). \end{align}\]

4. Arbitrary submodel sampling. Let \(\{B_1,\dots,B_m\}\) be a partition of \([b]\) into disjoint blocks of arbitrary indices, i.e., \[\begin{align} B_1 \cup \cdots \cup B_m = [b], \quad B_k \cap B_l = \emptyset \quad\textrm{for } k\neq l. \end{align}\] At each iteration, pick block \(B_i\) with probability \(p_i\) (where \(\sum_{i=1}^m p_i = 1\)) and set \(\hat{S} = B_i\).

   For \(i\in[b]\), let \(k(i)\) denote the unique block with \(i\in B_{k(i)}\), and let \(\underline{b}_k :=\min B_k\). Then \[\begin{align} F_i &= {\mathbb{P}}\left(\hat{s} \leq i\right) = \sum_{j: \underline{b}_j \leq i} p_j, \\ Q_i &= {\mathbb{P}}\left(i\in \hat{S}\right) = p_{k(i)}. \end{align}\] The expected cost per iteration is \[\begin{align} {\mathbb{E}}\left[\mathrm{cost}(\hat{S})\right] = c_{\mathrm{ov}} + \sum_{i=1}^b c_i \left( \sum_{k: \underline{b}_k \leq i} p_k \right) + \sum_{i=1}^b c_i^{\sharp} p_{k(i)}. \end{align}\]

We now consider the algorithm’s performance under the two smoothness regimes.

11.1 Smooth Case↩︎

According to 7, under 2, 1 run with stepsizes \(\gamma_i^k = \frac{1}{L_{i,S^k}^0}\) guarantees that \[\begin{align} \left( \min_{i\in[b]} w_i \right) \frac{1}{K} \sum_{k=0}^{K-1} \sum_{i=1}^b {\mathbb{E}}\left[\left\| \nabla_i f(X^k) \right\|^2_{(i) \star}\right] &\leq& \frac{1}{K} \sum_{k=0}^{K-1} \sum_{i=1}^b w_i {\mathbb{E}}\left[\left\| \nabla_i f(X^k) \right\|^2_{(i) \star}\right] \\ &\leq& \frac{f(X^0) - f^{\star}}{K}, \end{align}\] where \(w_i :={\mathbb{E}}\left[\frac{\mathbb{I}\left(i\in \hat{S}\right)}{2 L_{i,\hat{S}}^0}\right]\). Thus, to ensure that \(\frac{1}{K} \sum_{k=0}^{K-1} \sum_{i=1}^b {\mathbb{E}}\left[\left\| \nabla_i f(X^k) \right\|^2_{(i) \star}\right] \leq \varepsilon\), it suffices to run it for \[K = \left\lceil \frac{f(X^0) - f^{\star}}{\varepsilon \left( \min_{i\in[b]} w_i \right)} \right\rceil\] iterations. Now, recall from 16 that the expected cost of a single iteration is \[\begin{align} {\mathbb{E}}\left[\mathrm{cost}(S^k)\right] = c_{\mathrm{ov}} + \sum_{i=1}^b c_i {\mathbb{P}}\left(\min S^k \leq i\right) + \sum_{i=1}^b c_i^{\sharp} {\mathbb{P}}\left(i\in S^k\right). \end{align}\] Hence, the expected cost of the entire optimization procedure can be written as \[\begin{align} \mathrm{cost}_{\varepsilon}(\mathcal{D}) &=& K \times {\mathbb{E}}\left[\mathrm{cost}(\hat{S})\right] \\ &=& K \times \left( c_{\mathrm{ov}} + \sum_{i=1}^b c_i {\mathbb{P}}\left(\min \hat{S} \leq i\right) + \sum_{i=1}^b c_i^{\sharp} {\mathbb{P}}\left(i\in \hat{S}\right) \right) \\ &\propto& \frac{c_{\mathrm{ov}} + \sum_{i=1}^b c_i {\mathbb{P}}\left(\min \hat{S} \leq i\right) + \sum_{i=1}^b c_i^{\sharp} {\mathbb{P}}\left(i\in \hat{S}\right)}{\min_{i\in[b]} {\mathbb{E}}\left[\frac{\mathbb{I}\left(i\in \hat{S}\right)}{2 L_{i,\hat{S}}^0}\right]}, \end{align}\] and the cost minimization problem to be solved is \[\begin{align} \label{eq:cost95opt95problem} \min_{\mathcal{D}: \mathfrak{P}([b]) \to [0,1], \sum_{S\subseteq[b]} \mathcal{D}(S) = 1} \frac{c_{\mathrm{ov}} + \sum_{i=1}^b c_i {\mathbb{P}}\left(\min \hat{S} \leq i\right) + \sum_{i=1}^b c_i^{\sharp} {\mathbb{P}}\left(i\in \hat{S}\right)}{\min_{i\in[b]} {\mathbb{E}}\left[\frac{\mathbb{I}\left(i\in \hat{S}\right)}{2 L_{i,\hat{S}}^0}\right]}. \end{align}\tag{17}\]

The task above is an optimization over probability distributions on the power set \(\mathfrak{P}([b])\), which has dimension \(2^b\), making a direct solution intractable for large \(b\). Instead of tackling it in full generality, we can restrict \(\mathcal{D}\) to some parametric family. For certain such families, the ratio objective simplifies to a linear–fractional program in the cumulative marginals, which has a closed form solution or can be solved efficiently (e.g., via the Dinkelbach algorithm [42]).

Let us now consider some specific examples, starting with the procedure considered in the main part of this paper.

11.1.1 Randomized Progressive Training↩︎

We first consider the randomized progressive training setting introduced in 4. Under this sampling strategy, we have \[\begin{align} F_i = {\mathbb{P}}\left(\hat{s}\leq i\right) = \sum_{j=1}^i p_j, \qquad Q_i = {\mathbb{P}}\left(i\in \hat{S}\right) = {\mathbb{P}}\left(\hat{s}\leq i\right) = F_i, \end{align}\] and hence \[\begin{align} {\mathbb{E}}\left[\mathrm{cost}(\hat{S})\right] = c_{\mathrm{ov}} + \sum_{i=1}^b (c_i + c_i^{\sharp}) F_i = c_{\mathrm{ov}} + \sum_{i=1}^b (c_i + c_i^{\sharp}) \sum_{j=1}^i p_j. \end{align}\] Combining it with the fact that \[\begin{align} {\mathbb{E}}\left[\frac{\mathbb{I}\left(i\in \hat{S}\right)}{2 L_{i,\hat{S}}^0}\right] = \sum_{s=1}^b \frac{\mathbb{I}\left(i \in \{s,\dots,b\}\right)}{2 L_{i, \{s,\dots,b\}}^0} p_s = \sum_{s=1}^i \frac{p_s}{2 L_{i, \{s,\dots,b\}}^0}, \end{align}\] we get \[\begin{align} \label{eq:costK} \mathrm{cost}_{\varepsilon}(\mathcal{D}) \propto \frac{{\mathbb{E}}\left[\mathrm{cost}(\hat{S})\right]}{\min_{i\in[b]} {\mathbb{E}}\left[\frac{\mathbb{I}\left(i\in \hat{S}\right)}{2 L_{i,\hat{S}}^0}\right]} = \frac{c_{\mathrm{ov}} + \sum_{i=1}^b (c_i + c_i^{\sharp}) \sum_{j=1}^i p_j}{\min_{i\in[b]} \left\{ \sum_{s=1}^i \frac{p_s}{2 L_{i, \{s,\dots,b\}}^0} \right\}}. \end{align}\tag{18}\]

11.1.1.1 Optimal probabilities.

First, note that the numerator can be rewritten as \[\begin{align} c_{\mathrm{ov}} + \sum_{i=1}^b (c_i + c_i^{\sharp}) \sum_{j=1}^i p_j =c_{\mathrm{ov}}+\sum_{j=1}^b \left[\sum_{i\geq j} (c_i + c_i^{\sharp})\right] p_j =\sum_{j=1}^b \left[c_{\mathrm{ov}}+\sum_{i\geq j} (c_i + c_i^{\sharp})\right] p_j, \end{align}\] where the second equality follows from \(\sum_{j=1}^b p_j=1\). Denote \[d_j :=c_{\mathrm{ov}}+\sum_{i\geq j} (c_i + c_i^{\sharp})\] for \(j \in [b]\), and let \[\delta_{i,s}:=\frac{1}{2L_{i, \{s,\dots,b\}}}\] for \(i\in [b]\) and \(s\leq i\). Clearly \[\begin{align} \label{eq:orderofd} d_1 > d_2 > \ldots > d_b \end{align}\tag{19}\] and \[\begin{align} \label{eq:orderofdelta} \delta_{i,1} \leq \delta_{i,2} \leq \ldots \leq \delta_{i,i} \qquad \forall i\in [b]. \end{align}\tag{20}\] Based on 18 , the search for optimal probabilities reduces to solving the following linear fractional program: \[\label{eq:fl1} \begin{align} \min_{p, t} \quad & \frac{d_1p_1+\ldots+d_bp_b}{t} \\ \text{s.t.} \quad & p_1,\ldots,p_b\geq 0\\ & p_1+\ldots+p_b=1\\ & t\geq 0 \\ & t\leq \delta_{1,1}p_1 \\ & \vdots\\ & t\leq \delta_{i,1}p_1+\ldots+\delta_{i,i}p_i\\ & \vdots\\ & t\leq \delta_{b,1}p_1+\ldots+\delta_{b,b}p_b. \end{align}\tag{21}\] This program can be written equivalently as \[\label{eq:fl2} \begin{align} \min_{q} \quad & {d_1 q_1+\ldots+d_b q_b} \\ \text{s.t.} \quad & q_1,\ldots,q_b\geq 0\\ & \delta_{1,1}q_1\geq 1 \\ & \vdots\\ & \delta_{i,1}q_1+\ldots+\delta_{i,i}q_i\geq 1\\ & \vdots\\ & \delta_{b,1}q_1+\ldots+\delta_{b,b}q_b \geq 1. \end{align}\tag{22}\]

Based on that, we can derive a recursive prescription for the optimal probabilities.

First, we show that if \(\sum_{s=1}^i \delta_{i,s} q_s>1\) and \(q_i>0\) for some \(i\), then one can shift a small amount of mass from \(q_i\) to \(q_{i+1}\) and strictly reduce the objective, without violating any constraints.

Let us now use 2 to derive a recursive construction of the optimal probabilities. Let \(q^\star = (q_1^\star, \ldots, q_b^\star)\) be a solution of 22 . As established in 23 , any positive coordinate forces its constraint to be tight. Constraint \(1\) is \(\delta_{1,1} q_1^\star \geq 1\), meaning that \(q_1^\star > 0\), and so \(\delta_{1,1} q_1^\star = 1\). Thus \[\begin{align} q_1^\star = \frac{1}{\delta_{1,1}} = 2 L_{1, \{1,\dots,b\}}. \end{align}\] Now, suppose that \(q_1^\star,\dots,q_{i-1}^\star\) have already been determined. Then, constraint \(i\) reads \[\begin{align} \sum_{s=1}^{i-1} \delta_{i,s} q_s^\star + \delta_{i,i} q_i^\star \geq 1. \end{align}\] Define the residual \[\begin{align} r_i^\star :=1 - \sum_{s=1}^{i-1} \delta_{i,s} q_s^\star = 1 - \sum_{s=1}^{i-1} \frac{q_s^\star}{2L_{i, \{s,\dots,b\}}}. \end{align}\] If \(r_i^\star \le 0\), then \((q_1^\star,\dots,q_{i-1}^\star)\) already satisfy the constraint, so by 23 , we must have \(q_i^\star=0\). If \(r_i^\star >0\), then for \((q_1^\star,\dots,q_i^\star)\) to satisfy the constraint, we need \(q_i^\star>0\), and hence by 2, we have \(\sum_{s=1}^i \delta_{i,s} q_s^\star = 1\), meaning that \[\begin{align} q_i^\star = \frac{r_i^\star}{\delta_{i,i}} = 2 r_i^\star L_{i, \{i,\dots,b\}}. \end{align}\] Combining the above yields the recursion \[\begin{align} q_1^\star &= 2 L_{1, \{1,\dots,b\}}, \\ q_i^\star &= 2 [r_i^\star]_+ L_{i, \{i,\dots,b\}},\qquad r_i^\star = 1-\sum_{s=1}^{i-1} \frac{q_s^\star}{2L_{i, \{s,\dots,b\}}},\quad i\in\{2,\ldots,b\}, \end{align}\] where \([x]_+ :=\max\{x,0\}\). The optimal probabilities are finally recovered by normalization \[\begin{align} p_i^\star = \frac{q_i^\star}{\sum_{j=1}^b q_j^\star}. \end{align}\]

11.1.2 \(\tau\)-nice Sampling↩︎

Every \(\tau\)-subset has probability \(\binom{b}{\tau}^{-1}\). Thus \[\begin{align} {\mathbb{E}}\left[\frac{\mathbb{I}\left(i\in\hat{S}\right)}{2L^0_{i,\hat{S}}}\right] = \frac{1}{\binom{b}{\tau}} \sum_{\substack{S\subseteq[b]\\ |S|=\tau}} \frac{\mathbb{I}\left(i\in S\right)}{2L^0_{i,S}}, \end{align}\] and hence \[\begin{align} \mathrm{cost}_{\varepsilon}(\mathcal{D}) &\propto& \frac{{\mathbb{E}}\left[\mathrm{cost}(\hat{S})\right]}{\min_{i\in[b]} {\mathbb{E}}\left[\frac{\mathbb{I}\left(i\in \hat{S}\right)}{2 L_{i,\hat{S}}^0}\right]} = \frac{c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( 1 - \frac{\binom{b-j}{\tau}}{\binom{b}{\tau}} \right) + \frac{\tau}{b} \sum_{j=1}^b c_j^{\sharp}}{\min_{i\in[b]} \frac{1}{\binom{b}{\tau}} \sum_{\substack{S\subseteq[b]\\ |S|=\tau}} \frac{\mathbb{I}\left(i\in S\right)}{2L^0_{i,S}}} \\ &=& \frac{c_{\mathrm{ov}} \binom{b}{\tau} + \sum_{j=1}^b c_j \left( \binom{b}{\tau} - \binom{b-j}{\tau} \right) + \binom{b}{\tau} \frac{\tau}{b} \sum_{j=1}^b c_j^{\sharp}}{\min_{i\in[b]} \left\{ \sum_{\substack{S\subseteq[b]\\ |S|=\tau}} \frac{\mathbb{I}\left(i\in S\right)}{2L^0_{i,S}} \right\}}. \end{align}\] In general there is no closed-form expression for \(\tau\) minimizing this cost because the denominator depends on \(\{L^0_{i,S}\}_{S\subseteq[b]}\) in a highly problem-specific way. That said, it can be shown that \(\tau=b\) need not be optimal (indeed, it can be very sub-optimal). For \(\tau=b\), we always have \(S=[b]\), and the cost becomes \[\begin{align} \mathrm{cost}_{\varepsilon}(\mathcal{D}) &\propto& \frac{c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( 1 - \binom{b-j}{b} \right) + \sum_{j=1}^b c_j^{\sharp}}{\min_{i\in[b]} \frac{1}{2L^0_{i,[b]}}} \\ &=& 2 \max_{i\in[b]} L^0_{i,[b]} \left( c_{\mathrm{ov}} + \sum_{j=1}^b \left( c_j + c_j^{\sharp} \right) \right) \end{align}\] (recall that we use the convention \(\binom{n}{k}=0\) for \(n<k\)).

Now, let us make a simplifying assumption that \(L^0_{i,S}\) depends only on \(i\) and \(|S|\) and define \(L^0_{i,\tau} :=L^0_{i,S}\) for \(|S|=\tau\). Then \[\begin{align} \mathrm{cost}_{\varepsilon}(\mathcal{D}) &\propto& \frac{c_{\mathrm{ov}} \binom{b}{\tau} + \sum_{j=1}^b c_j \left( \binom{b}{\tau} - \binom{b-j}{\tau} \right) + \binom{b}{\tau} \frac{\tau}{b} \sum_{j=1}^b c_j^{\sharp}}{\min_{i\in[b]} \left\{ \sum_{\substack{S\subseteq[b]\\ |S|=\tau}} \frac{\mathbb{I}\left(i\in S\right)}{2L^0_{i,\tau}} \right\}} \\ &=& \frac{c_{\mathrm{ov}} \binom{b}{\tau} + \sum_{j=1}^b c_j \left( \binom{b}{\tau} - \binom{b-j}{\tau} \right) + \binom{b}{\tau} \frac{\tau}{b} \sum_{j=1}^b c_j^{\sharp}}{\min_{i\in[b]} \left\{ \frac{1}{2L^0_{i,\tau}} \binom{b-1}{\tau-1} \right\}} \\ &=& 2 \max_{i\in[b]} L^0_{i,\tau} \left( \frac{b}{\tau} c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( \frac{b}{\tau} - \frac{\binom{b-j}{\tau}}{\binom{b-1}{\tau-1}} \right) + \sum_{j=1}^b c_j^{\sharp} \right). \end{align}\] Define \[\begin{align} A(\tau) :=\max_{i\in[b]} L^0_{i,\tau}, \qquad B(\tau) :=\frac{b}{\tau} c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( \frac{b}{\tau} - \frac{\binom{b-j}{\tau}}{\binom{b-1}{\tau-1}} \right) + \sum_{j=1}^b c_j^{\sharp}. \end{align}\] Then, the objective function to be minimized is \[\begin{align} f(\tau) :=A(\tau) B(\tau). \end{align}\] By definition, \(\tau=b\) is optimal if and only if \(f(b) \leq f(\tau)\) for all \(\tau \in \{1,\dots,b-1\}\).

First, we show that \(B\) is decreasing in \(\tau\). To this end, let \(1 \leq \tau_1 < \tau_2 \leq b\) and consider the difference \[\begin{align} B(\tau_1) - B(\tau_2) &=& \left( \frac{b}{\tau_1} - \frac{b}{\tau_2} \right) c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( \frac{b}{\tau_1} - \frac{\binom{b-j}{\tau_1}}{\binom{b-1}{\tau_1-1}} - \frac{b}{\tau_2} + \frac{\binom{b-j}{\tau_2}}{\binom{b-1}{\tau_2-1}} \right) \\ &=& \left( \frac{b}{\tau_1} - \frac{b}{\tau_2} \right) c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( h_j(\tau_1) - h_j(\tau_2) \right), \end{align}\] where \(h_j(\tau) :=\frac{b}{\tau} - \frac{\binom{b-j}{\tau}}{\binom{b-1}{\tau-1}}\). Clearly, the first term is positive. Let us focus on the second term. Using Pascal’s identity, we have \[\begin{align} h_{j+1}(\tau)-h_j(\tau) &=\frac{\binom{b-j}{\tau}-\binom{b-j-1}{\tau}}{\binom{b-1}{\tau-1}} =\frac{\binom{b-j-1}{\tau-1}}{\binom{b-1}{\tau-1}}. \end{align}\] Therefore, \[\begin{align} h_j(\tau) = h_j(\tau) - h_0(\tau) = \sum_{m=0}^{j-1} \left( h_{m+1}(\tau)-h_m(\tau) \right) = \sum_{m=0}^{j-1} \frac{\binom{b-m-1}{\tau-1}}{\binom{b-1}{\tau-1}}, \end{align}\] where \[\begin{align} \frac{\binom{b-m-1}{\tau-1}}{\binom{b-1}{\tau-1}} = \frac{b - \tau}{b - m - \tau} \frac{\binom{b-m-1}{\tau}}{\binom{b-1}{\tau}} \geq \frac{\binom{b-m-1}{\tau}}{\binom{b-1}{\tau}}, \end{align}\] and the inequality is strict for \(m>0\). It follows that for any \(j \in [b]\) \[\begin{align} h_j(\tau) = \sum_{m=0}^{j-1} \frac{\binom{b-m-1}{\tau-1}}{\binom{b-1}{\tau-1}} \geq \sum_{m=0}^{j-1} \frac{\binom{b-m-1}{\tau}}{\binom{b-1}{\tau}} = h_j(\tau+1). \end{align}\] Thus, except for the trivial case when \(j=1\) (where \(h_1(\tau)\equiv1\)), the function \(h_j(\tau)\) is strictly decreasing in \(\tau\), implying that \(h_j(\tau_1) - h_j(\tau_2) > 0\), which proves that \(B(\tau_1) > B(\tau_2)\). Thus \[\begin{align} B(\tau) \geq B(b) \quad \forall \tau \leq b, \end{align}\] with strict inequality if \(c_{\mathrm{ov}} > 0\) and \(\tau < b\).

Now, note that in general larger \(\tau\) leads to a higher Lipschitz constant, and hence one may expect \(A\) to be an increasing function of \(\tau\). If \(A(b)\) is much larger than \(A(\tau)\), it may compensate for the decrease in \(B(\tau)\), resulting in \(f(\tau) = A(\tau) B(\tau) < A(b) B(b) = f(b)\), and so \(\tau=b\) may not be optimal.
As an example, suppose that \(L^0_{i,\tau}\) scales linearly with \(\tau\), i.e, \(A(\tau) = \max_{i\in[b]} \tau L^0_{i}\) for some \(L^0_{i} \geq 0\). Then \[\begin{align} \mathrm{cost}_{\varepsilon}(\mathcal{D}) &\propto& \max_{i\in[b]} L^0_{i} \left( b c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( b - \tau \frac{\binom{b-j}{\tau}}{\binom{b-1}{\tau-1}} \right) + \tau \sum_{j=1}^b c_j^{\sharp} \right) \\ &=& \max_{i\in[b]} L^0_{i} \left( b c_{\mathrm{ov}} + b \sum_{j=1}^b c_j + \tau \sum_{j=1}^b \left( c_j^{\sharp} - c_j \frac{\binom{b-j}{\tau}}{\binom{b-1}{\tau-1}} \right) \right). \end{align}\] Now, consider the function \[\begin{align} \Phi(\tau) :=\tau \sum_{j=1}^b \left( c_j^{\sharp} - c_j \frac{\binom{b-j}{\tau}}{\binom{b-1}{\tau-1}} \right). \end{align}\] Note that \[\begin{align} \Phi(\tau+1) - \Phi(\tau) &=& (\tau+1) \sum_{j=1}^b \left( c_j^{\sharp} - c_j \frac{\binom{b-j}{\tau+1}}{\binom{b-1}{\tau}} \right) - \tau \sum_{j=1}^b \left( c_j^{\sharp} - c_j \frac{\binom{b-j}{\tau}}{\binom{b-1}{\tau-1}} \right) \\ &=& \sum_{j=1}^b c_j \left( \tau \frac{\binom{b-j}{\tau}}{\binom{b-1}{\tau-1}} - (\tau+1) \frac{\binom{b-j}{\tau+1}}{\binom{b-1}{\tau}} \right) + \sum_{j=1}^b c_j^{\sharp} \\ &=& \sum_{j=1}^b c_j \frac{j (b-j)! (b-\tau-1)!}{(b-1)! (b-j-\tau)!} + \sum_{j=1}^b c_j^{\sharp} \geq 0 \end{align}\] (with the convention that the right-hand side is \(0\) when \(b-j-\tau<0\)). Moreover, the increment is strictly positive if either \(\sum_{j=1}^b c_j^{\sharp}>0\) or there exists \(j\) with \(c_j>0\) and \(b-j-\tau\geq 0\). Hence \(\Phi(\tau)\) is non-decreasing in \(\tau\), and strictly increasing whenever one of these conditions holds, and the optimal choice is \(\tau^{\star}=1\).

11.1.3 \(\tau\)-Submodel Sampling↩︎

The denominator is \[\begin{align} {\mathbb{E}}\left[\frac{\mathbb{I}\left(i\in\hat{S}\right)}{2L^0_{i,\hat{S}}}\right] = \sum_{j=1}^{b-\tau+1} p_j \frac{\mathbb{I}\left(i\in\{j,\dots,j+\tau-1\}\right)}{2L^0_{i,\{j,\dots,j+\tau-1\}}} = \sum_{j=\max\{1,i-\tau+1\}}^{\min\{i,b-\tau+1\}} \frac{p_j}{2L^0_{i,\{j,\dots,j+\tau-1\}}}. \end{align}\] Hence the total expected cost is proportional to \[\begin{align} \mathrm{cost}_{\varepsilon}(\mathcal{D}) &\propto& \frac{{\mathbb{E}}\left[\mathrm{cost}(\hat{S})\right]}{\min_{i\in[b]} {\mathbb{E}}\left[\frac{\mathbb{I}\left(i\in\hat{S}\right)}{2L^0_{i,\hat{S}}}\right]} \\ &=& \frac{c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left(\sum_{i=1}^{\min\{j,b-\tau+1\}} p_i\right) + \sum_{j=1}^b c_j^{\sharp} \left(\sum_{i=\max\{1,j-\tau+1\}}^{\min\{j,b-\tau+1\}} p_i\right)}{\min_{i\in[b]} \left\{ \sum_{j=\max\{1,i-\tau+1\}}^{\min\{i,b-\tau+1\}} \frac{p_j}{2L^0_{i,\{j,\dots,j+\tau-1\}}} \right\}}. \end{align}\]

11.1.3.1 Partitioned \(\tau\)-submodel sampling.

Let us consider a special case of the above sampling scheme where the submodels assigned non-zero probability partition \([b]\). For simplicity, suppose that \(b\) is divisible by \(\tau\), let \(m=\frac{b}{\tau}\), and define the block start indices via \[\begin{align} s_k=(k-1)\tau+1,\qquad k=1,\dots,m. \end{align}\] The algorithm then picks block \(B_k :=\{s_k, \dots, s_k+\tau-1\}\) with probability \(p_{s_k}> 0\) (where \(\sum_{k=1}^m p_{s_k} = 1\)). This is equivalent to the submodel sampling with starting-index distribution satisfying \[\begin{align} p_{s_k}=p_{(k-1)\tau+1} > 0 \quad (k=1,\dots,m), \qquad p_j=0 \text{ otherwise}. \end{align}\] Plugging this choice into the general submodel expressions immediately gives \[\begin{align} F_i &=& \sum_{j=1}^{\min\{i, b-\tau+1\}} p_j = \sum_{k=1}^{\min\{m, \left\lfloor \frac{(i-1)}{\tau} \right\rfloor + 1\}} p_{s_k}, \\ Q_i &=& \sum_{j=\max\{1, i-\tau+1\}}^{\min\{i, b-\tau+1\}} p_j = p_{s_{\left\lceil \frac{i}{\tau} \right\rceil}}, \\ {\mathbb{E}}\left[\frac{\mathbb{I}\left(i\in\hat{S}\right)}{2L^0_{i,\hat{S}}}\right] &=& \sum_{j=\max\{1,i-\tau+1\}}^{\min\{i,b-\tau+1\}} \frac{p_j}{2L^0_{i,\{j,\dots,j+\tau-1\}}} = \frac{p_{s_{\left\lceil \frac{i}{\tau} \right\rceil}}}{2L^0_{i,B_{\left\lceil \frac{i}{\tau} \right\rceil}}}. \end{align}\] Therefore the full cost reduces to \[\begin{align} \label{eq:part95tau95submodel95cost} \mathrm{cost}_{\varepsilon}(\mathcal{D}) &\propto& \frac{c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( \sum_{k=1}^{\min\{m, \left\lfloor \frac{(j-1)}{\tau} \right\rfloor + 1\}} p_{s_k} \right) + \sum_{j=1}^b c_j^{\sharp} p_{s_{\left\lceil \frac{j}{\tau} \right\rceil}}}{\min_{i\in[b]} \left\{ \frac{p_{s_{\left\lceil \frac{i}{\tau} \right\rceil}}}{2L^0_{i,B_{\left\lceil \frac{i}{\tau} \right\rceil}}} \right\}} \nonumber \\ &=& \frac{c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( \sum_{k=1}^{\left\lceil \frac{j}{\tau} \right\rceil} p_{s_k} \right) + \sum_{j=1}^b c_j^{\sharp} p_{s_{\left\lceil \frac{j}{\tau} \right\rceil}}}{\min_{i\in[b]} \left\{ \frac{p_{s_{\left\lceil \frac{i}{\tau} \right\rceil}}}{2L^0_{i,B_{\left\lceil \frac{i}{\tau} \right\rceil}}} \right\}}. \end{align}\tag{24}\] Now, note that \[\begin{align} \sum_{j=1}^b c_j \left( \sum_{k=1}^{\left\lceil \frac{j}{\tau} \right\rceil} p_{s_k} \right) &= \sum_{j=1}^\tau c_j p_{s_1} + \sum_{j=\tau+1}^{2\tau} c_j \left( p_{s_1} + p_{s_2} \right) + \ldots + \sum_{j=(m-1)\tau+1}^{m\tau} c_j \left( p_{s_1} + \ldots + p_{s_m} \right) \\ &= p_{s_1} \sum_{j=1}^\tau c_j + \left( p_{s_1} + p_{s_2} \right) \sum_{j=\tau+1}^{2\tau} c_j + \ldots + \left( p_{s_1} + \ldots + p_{s_m} \right) \sum_{j=(m-1)\tau+1}^{m\tau} c_j \\ &= p_{s_1} \sum_{j=1}^{m \tau} c_j + p_{s_2} \sum_{j=\tau+1}^{m\tau} c_j + \ldots + p_{s_m} \sum_{j=(m-1)\tau+1}^{m\tau} c_j \end{align}\] and \[\begin{align} \sum_{j=1}^b c_j^{\sharp} p_{s_{\left\lceil \frac{j}{\tau} \right\rceil}} = \sum_{j=1}^\tau c_j^{\sharp} p_{s_1} + \sum_{j=\tau+1}^{2\tau} c_j^{\sharp} p_{s_2} + \ldots + \sum_{j=(m-1)\tau+1}^{m\tau} c_j^{\sharp} p_{s_m}. \end{align}\] Hence \[\begin{align} &\mathrm{cost}_{\varepsilon}(\mathcal{D}) \\ &\propto \frac{1}{\min_{i\in[b]} \left\{ \frac{p_{s_{\left\lceil \frac{i}{\tau} \right\rceil}}}{2L^0_{i,B_{\left\lceil \frac{i}{\tau} \right\rceil}}} \right\}} \left[c_{\mathrm{ov}} + p_{s_1} \left( \sum_{j=1}^{m \tau} c_j + \sum_{j=1}^\tau c_j^{\sharp} \right) + p_{s_2} \left( \sum_{j=\tau+1}^{m\tau} c_j + \sum_{j=\tau+1}^{2\tau} c_j^{\sharp} \right) + \ldots \right. \\ &\qquad\qquad\qquad\qquad\qquad\qquad\left.+ p_{s_m} \left( \sum_{j=(m-1)\tau+1}^{m\tau} c_j + \sum_{j=(m-1)\tau+1}^{m\tau} c_j^{\sharp} \right)\right] \\ &= \frac{p_{s_1} d_1 + p_{s_2} d_2 + \ldots + p_{s_m} d_m}{\min_{i\in[b]} \left\{ \frac{p_{s_{\left\lceil \frac{i}{\tau} \right\rceil}}}{2L^0_{i,B_{\left\lceil \frac{i}{\tau} \right\rceil}}} \right\}}, \end{align}\] where we used the fact that \(\sum_{j=1}^{m} p_{s_j} = 1\) and denoted \(d_i :=c_{\mathrm{ov}} + \sum_{j=(i-1)\tau+1}^{m \tau} c_j + \sum_{j=(i-1)\tau+1}^{i \tau} c_j^{\sharp}\).

11.1.3.2 Optimal probabilities for fixed \(\tau\).

We follow an approach similar to that in 11.1.1 to find the optimal probabilities for a fixed choice of \(\tau\). For \(i\in B_k\), define \[\begin{align} \delta_{i,k}:=\frac{1}{2L^0_{i,B_k}} \end{align}\] and set \[\begin{align} \delta_k^{\min}:=\min_{i\in B_k}\delta_{i,k} = \frac{1}{2\max_{i\in B_k} L^0_{i,B_k}}. \end{align}\] Then the expected cost can be represented as \[\begin{align} \mathrm{cost}_{\varepsilon}(\mathcal{D}) \propto \frac{d_1 p_{s_1} + \cdots + d_m p_{s_m}}{\min_{i\in[b]} \left\{ \frac{p_{s_{\left\lceil \frac{i}{\tau} \right\rceil}}}{2L^0_{i,B_{\left\lceil \frac{i}{\tau} \right\rceil}}} \right\}} = \frac{d_1 p_{s_1} + \cdots + d_m p_{s_m}}{\min_{k\in[m]} \min_{i\in B_k} \{\delta_{i,k} p_{s_k}\}}. \end{align}\] This corresponds to the linear fractional program \[\label{eq:partitioned95submod95lfp} \begin{align} \min_{p,t}\quad & \frac{d_1 p_{s_1} + \cdots + d_m p_{s_m}}{t} \\ \text{s.t.}\quad & p_{s_1}, \ldots, p_{s_m} \geq 0, \\ & p_{s_1}+\cdots+p_{s_m} = 1, \\ & t \geq 0, \\ & t \leq \delta_{i,k} p_{s_k}, \qquad k\in[m], i\in B_k. \end{align}\tag{25}\] The standard change of variables \(q_k = p_{s_k} / t\) yields the equivalent linear program \[\label{eq:fl295partitioned} \begin{align} \min_{q} \quad & d_1 q_1 + \cdots + d_m q_m \\ \text{s.t.} \quad & q_1, \ldots, q_m \geq 0, \\ & \delta_{i,k} q_k \geq 1 \qquad k\in[m], i\in B_k. \end{align}\tag{26}\] Because each block \(k\) only appears in constraints of the form \(\delta_{i,k} q_k \geq 1\) (for \(i\in B_k\)), the constraints for block \(k\) reduce to the single constraint \[\begin{align} \delta_k^{\min} q_k \geq 1 \quad\Longleftrightarrow\quad q_k \geq \frac{1}{\delta_k^{\min}}. \end{align}\] Hence 26 is separable and its optimal solution is \[\begin{align} q_k^\star = \frac{1}{\delta_k^{\min}}, \qquad k\in[m], \end{align}\] with optimal objective value \[\begin{align} \sum_{k=1}^m d_k q_k^\star = \sum_{k=1}^m \frac{d_k}{\delta_k^{\min}}. \end{align}\]

Now, let us introduce dual variables \(\lambda_i\ge0\) for the constraints \(\delta_{i,k} q_k \geq 1\), \(i\in B_k\). The dual of 26 is \[\begin{align} \max_{\lambda}\quad & \lambda_1 + \ldots + \lambda_b \\ \text{s.t.}\quad & \lambda_1, \ldots, \lambda_b \geq 0, \\ & d_k \geq \sum_{i\in B_k} \lambda_i \delta_{i,k}, \qquad k\in[m]. \end{align}\] To certify optimality of \(q^\star\), for each block \(k\), choose an index \(i_k\in B_k\) attaining the minimum \(\delta_k^{\min}\) and set \[\begin{align} \lambda_{i_k} = \frac{d_k}{\delta_{i_k,k}}, \qquad \lambda_i=0 \text{ for } i\not\in\{i_1,\dots,i_m\}. \end{align}\] Then for each block \(k\), \[\begin{align} \sum_{i\in B_k} \delta_{i,k} \lambda_i = \delta_{i_k,k} \lambda_{i_k} = d_k, \end{align}\] so the dual constraints hold with equality and the dual objective equals \[\begin{align} \sum_{i=1}^b \lambda_i = \sum_{k=1}^m \frac{d_k}{\delta_k^{\min}}, \end{align}\] matching the primal objective. By strong duality \(q^\star\) is optimal.

From \(q_k^\star=\frac{1}{\delta_k^{\min}}\) and \(p_{s_k} = t q_k = q_k / \sum_{l=1}^m q_l\) we obtain the optimal block probabilities \[\begin{align} p_{s_k}^\star = \frac{\frac{1}{\delta_k^{\min}}}{\sum_{l=1}^m \frac{1}{\delta_l^{\min}}} = \frac{\max_{i\in B_k} L^0_{i,B_k}}{\sum_{l=1}^m \max_{i\in B_l} L^0_{i,B_l}}. \end{align}\] so that each block’s probability is proportional to the worst-case local smoothness constant inside that block. For this choice, the minimal expected cost is \[\begin{align} \mathrm{cost}_{\varepsilon}(\mathcal{D}) \propto \sum_{k=1}^m \frac{d_k}{\delta_k^{\min}} = 2\sum_{k=1}^m d_k \max_{i\in B_k} L^0_{i,B_k}. \end{align}\]

11.1.3.3 Choosing \(\tau\).

We now show that the cost is not necessarily minimized by choosing \(\tau = b\). To this end, we want to minimize the function \[\begin{align} \Phi(\tau) :=2\sum_{k=1}^{m} d_k(\tau) \max_{i\in B_k} L^0_{i,B_k} \end{align}\] (where we explicitly emphasize the dependence of \(d_k\) on \(\tau\)).

To gain intuition about which extreme (\(\tau = 1\) or \(\tau = b\)) may be preferable, assume that the costs are constant, i.e., \[\begin{align} c_j \equiv c, \qquad c_j^\sharp \equiv c^\sharp, \end{align}\] in which case \[\begin{align} \sum_{k=1}^m d_k(\tau) &= \sum_{k=1}^m \left( c_{\mathrm{ov}} + \sum_{j=(k-1)\tau+1}^{m \tau} c + \sum_{j=(k-1)\tau+1}^{k\tau} c^{\sharp} \right) \\ &= \sum_{k=1}^m \left( c_{\mathrm{ov}} + c \left( m \tau - (k-1) \tau \right) + c^{\sharp} \left( k\tau - (k-1)\tau \right) \right) \\ &= m c_{\mathrm{ov}} + c \tau \frac{m(m+1)}{2} + m \tau c^\sharp. \end{align}\] Suppose in addition that the worst-case local smoothness per block does not depend on the block index, i.e., \[\begin{align} \max_{i\in B_k} L^0_{i,B_k} \equiv L^0(\tau) \quad \forall k\in[m] \end{align}\] for some non-decreasing function \(L\). Then \[\begin{align} \Phi(\tau) = 2 L^0(\tau)\sum_{k=1}^{m} d_k(\tau) = 2 L^0(\tau) \left( \frac{b}{\tau} \left( c_{\mathrm{ov}}+\frac{bc}{2} \right) + \frac{bc}{2} + b c^\sharp \right). \end{align}\] Thus, the \(\tau\)-dependence of \(\Phi\) is the product of a non-decreasing factor \(L_0(\tau)\) and a factor that decreases like \(1/\tau\) plus additive constants. Consequently:

• If \(L^0(\tau)\) is constant in \(\tau\) (no worsening with larger blocks), then \(\Phi(\tau)\) is strictly decreasing in \(\tau\) and the minimizer is \(\tau^\star = b\).

• If \(L^0(\tau)\) grows sublinearly, then \(L^0(\tau) \frac{b}{\tau} \left( c_{\mathrm{ov}}+\frac{bc}{2} \right)\) is decreasing in \(\tau\), while \(L^0(\tau) \left( \frac{c b}{2} + b c^\sharp \right)\) is increasing. Hence, the optimal \(\tau\) may lie strictly between \(1\) and \(b\), depending on the relative magnitudes of the costs.

• If \(L^0(\tau)\) increases at least linearly in \(\tau\), then \(\Phi(\tau)\) is increasing in \(\tau\), and hence \(\tau^\star = 1\).

11.1.4 Arbitrary Submodel Sampling↩︎

First, note that the cost can be expressed block-wise as \[\begin{align} {\mathbb{E}}\left[\mathrm{cost}(\hat{S})\right] &= c_{\mathrm{ov}} + \sum_{i=1}^b c_i \left( \sum_{k:\underline{b}_k \leq i} p_k \right) + \sum_{i=1}^b c_i^{\sharp} p_{k(i)} \\ &= c_{\mathrm{ov}} + \sum_{k=1}^m \sum_{i \geq \underline{b}_k} c_i p_k + \sum_{k=1}^m \sum_{i\in B_k} c_i^{\sharp} p_{k} \\ &= \sum_{k=1}^m d_k p_k, \end{align}\] where \(d_k :=c_{\mathrm{ov}} + \sum_{j\geq \underline{b}_k} c_j + \sum_{j\in B_k} c_j^{\sharp}\). We also have \[\begin{align} {\mathbb{E}}\left[\frac{\mathbb{I}\left(i\in \hat{S}\right)}{2 L^0_{i,\hat{S}}}\right] = \frac{p_{k(i)}}{2 L^0_{i,B_{k(i)}}}, \end{align}\] and hence \[\begin{align} \mathrm{cost}_{\varepsilon}(\mathcal{D}) \propto \frac{{\mathbb{E}}\left[\mathrm{cost}(\hat{S})\right]}{\min_{i\in[b]} {\mathbb{E}}\left[\frac{\mathbb{I}\left(i\in\hat{S}\right)}{2L^0_{i,\hat{S}}}\right]} = \frac{\sum_{k=1}^m d_k p_k}{\min_{k\in[m]} \min_{i\in B_k} \left\{ \frac{p_k}{2 L^0_{i,B_k}} \right\}}. \end{align}\] Now, define \[\begin{align} \delta_{i,k} :=\tfrac{1}{2 L^0_{i,B_k}}, \qquad \delta_k^{\min} :=\min_{i\in B_k} \delta_{i,k}. \end{align}\] By the same linear fractional reduction as in 25 , the unique optimal solution is \[\begin{align} p_k^\star = \frac{\frac{1}{\delta_k^{\min}}}{\sum_{l=1}^m \frac{1}{\delta_l^{\min}}} = \frac{\max_{i\in B_k} L^0_{i,B_k}}{\sum_{l=1}^m \max_{i\in B_l} L^0_{i,B_l}}. \end{align}\] With this choice, the objective is \[\begin{align} \label{eq:cost95arbitrary95submod} \mathrm{cost}_{\varepsilon}(\mathcal{D}) &\propto& 2 \sum_{k=1}^m d_k \max_{i\in B_k} L^0_{i,B_k} = 2\sum_{k=1}^m \left( c_{\mathrm{ov}} + \sum_{j\geq \underline{b}_k} c_j + \sum_{j\in B_k} c_j^{\sharp} \right) \max_{i\in B_k} L^0_{i,B_k}. \end{align}\tag{27}\]

Let us compare it with the cost for full model training \[\begin{align} {\mathbb{E}}\left[\mathrm{cost}_{\textrm{full}}(K)\right] \propto 2 \max_{i\in[b]} L^0_{i,[b]} \left( c_{\mathrm{ov}} + \sum_{j=1}^b \left( c_j + c_j^{\sharp} \right) \right). \end{align}\]

11.2 Generalized Smooth Case↩︎

According to 10, under 3, 1 run with stepsizes \(\gamma_i^k = \big(L_{i,S^k}^0 + L_{i,S^k}^1 \left\| \nabla_i f(X^k) \right\|_{(i) \star}\big)^{-1}\) guarantees that \[\begin{align} \label{eq:avndion} \min_{k=0,\ldots,K-1} \sum_{i=1}^b \left[\frac{w_i}{\frac{1}{b} \sum_{l=1}^b w_l} {\mathbb{E}}\left[\left\| \nabla _i f(X^k) \right\|_{(i) \star}\right]\right] \leq \underline{\varepsilon}, \end{align}\tag{32}\] after \[\begin{align} \label{eq:oinqvr} K &=& \left\lceil \frac{2 \delta^0 \sum_{i=1}^b \frac{{\mathbb{P}}\left(i\in \hat{S}\right) {\mathbb{E}}\left[\left.L^0_{i,\hat{S}}\right\vert i\in \hat{S}\right]}{\left( {\mathbb{E}}\left[\left.L^1_{i,\hat{S}}\right\vert i\in \hat{S}\right] \right)^2}}{\underline{\varepsilon}^2 \left( \frac{1}{b} \sum_{l=1}^b \frac{{\mathbb{P}}\left(l\in \hat{S}\right)}{{\mathbb{E}}\left[\left.L^1_{l,\hat{S}}\right\vert l\in \hat{S}\right]} \right)^2} + \frac{2 \delta^0}{\underline{\varepsilon} \left( \frac{1}{b} \sum_{l=1}^b \frac{{\mathbb{P}}\left(l \in \hat{S}\right)}{{\mathbb{E}}\left[\left.L^1_{l,\hat{S}}\right\vert l\in \hat{S}\right]} \right)} \right\rceil \end{align}\tag{33}\] iterations, where \(\delta^0 :=f(X^0) - f^{\star}\) and \(w_i :=\frac{{\mathbb{P}}\left(i\in \hat{S}\right)}{{\mathbb{E}}\left[\left.L^1_{i,\hat{S}}\right\vert i\in \hat{S}\right]}\). To obtain a guarantee on an unweighted gradient sum, note that \[\begin{align} \underline{\varepsilon} &\geq& \min_{k=0,\ldots,K-1} \sum_{i=1}^b \left[\frac{w_i}{\frac{1}{b} \sum_{l=1}^b w_l} {\mathbb{E}}\left[\left\| \nabla _i f(X^k) \right\|_{(i) \star}\right]\right] \\ &\geq& \frac{\min_{i\in[b]} w_i}{\frac{1}{b} \sum_{l=1}^b w_l} \min_{k=0,\ldots,K-1} \sum_{i=1}^b {\mathbb{E}}\left[\left\| \nabla _i f(X^k) \right\|_{(i) \star}\right], \end{align}\] and hence, substituting in 32 and 33 , we have \[\begin{align} \min_{k=0,\ldots,K-1} \sum_{i=1}^b {\mathbb{E}}\left[\left\| \nabla _i f(X^k) \right\|_{(i) \star}\right] \leq \frac{\underline{\varepsilon} \frac{1}{b} \sum_{l=1}^b w_l}{\min_{i\in[b]} w_i} :=\varepsilon \end{align}\] after \[\begin{align} K &=& \left\lceil \frac{2 \delta^0 \sum_{i=1}^b \frac{{\mathbb{P}}\left(i\in \hat{S}\right) {\mathbb{E}}\left[\left.L^0_{i,\hat{S}}\right\vert i\in \hat{S}\right]}{\left( {\mathbb{E}}\left[\left.L^1_{i,\hat{S}}\right\vert i\in \hat{S}\right] \right)^2}}{\varepsilon^2 \min_{i\in[b]} \left( \frac{{\mathbb{P}}\left(i\in \hat{S}\right)}{{\mathbb{E}}\left[\left.L^1_{i,\hat{S}}\right\vert i\in \hat{S}\right]} \right)^2} + \frac{2 \delta^0}{\varepsilon \min_{i\in[b]} \frac{{\mathbb{P}}\left(i\in \hat{S}\right)}{{\mathbb{E}}\left[\left.L^1_{i,\hat{S}}\right\vert i\in \hat{S}\right]}} \right\rceil \end{align}\] iterations. Moreover, recall from 16 that the expected cost of a single iteration is \[\begin{align} {\mathbb{E}}\left[\mathrm{cost}(\hat{S})\right] = c_{\mathrm{ov}} + \sum_{i=1}^b c_i F_i + \sum_{i=1}^b c_i^{\sharp} Q_i. \end{align}\] Hence, using the fact that \[\begin{align} {\mathbb{E}}\left[\left.L^{\alpha}_{i,\hat{S}}\right\vert i\in \hat{S}\right] = \frac{{\mathbb{E}}\left[L^{\alpha}_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right]}{{\mathbb{P}}\left(i\in \hat{S}\right)} \end{align}\] for \(\alpha \in \{0,1\}\), the expected cost of the entire optimization procedure is \[\begin{align} &\mathrm{cost}_{\varepsilon}(\mathcal{D}) = K \times {\mathbb{E}}\left[\mathrm{cost}(\hat{S})\right] \\ &= \left\lceil \frac{2 \delta^0 \sum_{i=1}^b \frac{{\mathbb{P}}\left(i\in \hat{S}\right) {\mathbb{E}}\left[\left.L^0_{i,\hat{S}}\right\vert i\in \hat{S}\right]}{\left( {\mathbb{E}}\left[\left.L^1_{i,\hat{S}}\right\vert i\in \hat{S}\right] \right)^2}}{\varepsilon^2 \min_{i\in[b]} \left( \frac{{\mathbb{P}}\left(i\in \hat{S}\right)}{{\mathbb{E}}\left[\left.L^1_{i,\hat{S}}\right\vert i\in \hat{S}\right]} \right)^2} + \frac{2 \delta^0}{\varepsilon \min_{i\in[b]} \frac{{\mathbb{P}}\left(i\in \hat{S}\right)}{{\mathbb{E}}\left[\left.L^1_{i,\hat{S}}\right\vert i\in \hat{S}\right]}} \right\rceil \left( c_{\mathrm{ov}} + \sum_{i=1}^b c_i F_i + \sum_{i=1}^b c_i^{\sharp} Q_i \right) \\ &= \left\lceil \frac{2 \delta^0 \sum_{i=1}^b \frac{{\mathbb{P}}\left(i\in \hat{S}\right)^2 {\mathbb{E}}\left[L^0_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right]}{\left( {\mathbb{E}}\left[L^1_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right] \right)^2}}{\varepsilon^2 \min_{i\in[b]} \left( \frac{{\mathbb{P}}\left(i\in \hat{S}\right)^2}{{\mathbb{E}}\left[L^1_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right]} \right)^2} + \frac{2 \delta^0}{\varepsilon \min_{i\in[b]} \frac{{\mathbb{P}}\left(i\in \hat{S}\right)^2}{{\mathbb{E}}\left[L^1_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right]}} \right\rceil \left( c_{\mathrm{ov}} + \sum_{i=1}^b c_i F_i + \sum_{i=1}^b c_i^{\sharp} Q_i \right). \end{align}\] We will consider two regimes.

(1) The \(\mathcal{O}\left( \frac{1}{\varepsilon^2} \right)\) term dominates. Then the problem to be solved is \[\begin{align} \label{eq:cost95opt95prob95e2} \min_{\mathcal{D}: \mathfrak{P}([b]) \to [0,1], \sum_{S\subseteq[b]} \mathcal{D}(S) = 1} \underbrace{\frac{\sum_{i=1}^b \frac{{\mathbb{P}}\left(i\in \hat{S}\right)^2 {\mathbb{E}}\left[L^0_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right]}{\left( {\mathbb{E}}\left[L^1_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right] \right)^2}}{\min_{i\in[b]} \left( \frac{{\mathbb{P}}\left(i\in \hat{S}\right)^2}{{\mathbb{E}}\left[L^1_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right]} \right)^2} \left( c_{\mathrm{ov}} + \sum_{i=1}^b c_i F_i + \sum_{i=1}^b c_i^{\sharp} Q_i \right)}_{\propto \, {\mathbb{E}}\left[\mathrm{cost}_{\varepsilon^2}(K)\right]}. \end{align}\tag{34}\]

(2) The \(\mathcal{O}\left( \frac{1}{\varepsilon} \right)\) term dominates. Then the problem to be solved is \[\begin{align} \label{eq:cost95opt95prob95e} \min_{\mathcal{D}: \mathfrak{P}([b]) \to [0,1], \sum_{S\subseteq[b]} \mathcal{D}(S) = 1} \underbrace{\frac{1}{\min_{i\in[b]} \frac{{\mathbb{P}}\left(i\in \hat{S}\right)^2}{{\mathbb{E}}\left[L^1_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right]}} \left( c_{\mathrm{ov}} + \sum_{i=1}^b c_i F_i + \sum_{i=1}^b c_i^{\sharp} Q_i \right)}_{\propto \, {\mathbb{E}}\left[\mathrm{cost}_{\varepsilon}(K)\right]}. \end{align}\tag{35}\] As in 11.1, both tasks above involve optimization over probability distributions, which is intractable in general. Again, for certain parametric families, the objective simplifies to a linear-fractional program, which can be solved efficiently. Let us consider some specific examples.

11.2.1 Randomized Progressive Training↩︎

In this case, \(\hat{S} = \{s,\ldots,b\}\) with probability \(p_s\) for all \(s\in[b]\), and hence \[\begin{align} {\mathbb{P}}\left(i\in \hat{S}\right) = {\mathbb{P}}\left(s \leq i\right) = \sum_{s=1}^i p_s, \end{align}\] and \[\begin{align} {\mathbb{E}}\left[L^{\alpha}_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right] = \sum_{s=1}^i p_s L^{\alpha}_{i, \{s,\dots,b\}} \end{align}\] for \(\alpha\in\{0,1\}\). Moreover, following the same steps as in 11.1.1, \[\begin{align} {\mathbb{E}}\left[\mathrm{cost}(\hat{S})\right] = c_{\mathrm{ov}} + \sum_{j=1}^b (c_j + c_j^{\sharp}) \sum_{i=1}^j p_i = \sum_{i=1}^b \left[c_{\mathrm{ov}}+\sum_{j\geq i} (c_j + c_j^{\sharp})\right] p_i. \end{align}\] Hence, substituting into 34 and 35 , the respective optimization problems reduce to minimizing \[\begin{align} {\mathbb{E}}\left[\mathrm{cost}_{\varepsilon^2}(K)\right] &\propto& \frac{\sum_{i=1}^b \frac{{\mathbb{P}}\left(i\in \hat{S}\right)^2 {\mathbb{E}}\left[L^0_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right]}{\left( {\mathbb{E}}\left[L^1_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right] \right)^2}}{\min_{i\in[b]} \left( \frac{{\mathbb{P}}\left(i\in \hat{S}\right)^2}{{\mathbb{E}}\left[L^1_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right]} \right)^2} \left( \sum_{i=1}^b \left[c_{\mathrm{ov}}+\sum_{j\geq i} (c_j + c_j^{\sharp})\right] p_i \right) \\ &=& \frac{\sum_{i=1}^b \frac{\left( \sum_{s=1}^i p_s \right)^2 \sum_{s=1}^i p_s L^0_{i, \{s,\dots,b\}}}{\left( \sum_{s=1}^i p_s L^1_{i, \{s,\dots,b\}} \right)^2}}{\min_{i\in[b]} \left( \frac{\left( \sum_{s=1}^i p_s \right)^2}{\sum_{s=1}^i p_s L^1_{i, \{s,\dots,b\}}} \right)^2} \left( \sum_{i=1}^b \left[c_{\mathrm{ov}}+\sum_{j\geq i} (c_j + c_j^{\sharp})\right] p_i \right) \end{align}\] and \[\begin{align} {\mathbb{E}}\left[\mathrm{cost}_{\varepsilon}(K)\right] &\propto& \frac{1}{\min_{i\in[b]} \frac{{\mathbb{P}}\left(i\in \hat{S}\right)^2}{{\mathbb{E}}\left[L^1_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right]}} \left( \sum_{i=1}^b \left[c_{\mathrm{ov}}+\sum_{j\geq i} (c_j + c_j^{\sharp})\right] p_i \right) \\ &=& \frac{1}{\min_{i\in[b]} \frac{\left( \sum_{s=1}^i p_s \right)^2}{\sum_{s=1}^i p_s L^1_{i, \{s,\dots,b\}}}} \left( \sum_{i=1}^b \left[c_{\mathrm{ov}}+\sum_{j\geq i} (c_j + c_j^{\sharp})\right] p_i \right). \end{align}\]

Let us first focus on \({\mathbb{E}}\left[\mathrm{cost}_{\varepsilon}(K)\right]\), following a similar reasoning to that in 11.1.1. Denote \(d_i :=c_{\mathrm{ov}}+\sum_{j\geq i} (c_j + c_j^{\sharp})\) for \(i\in [b]\) and consider the linear fractional program \[\label{eq:fl1l0l1} \begin{align} \min_{p, t} \quad & \frac{d_1 p_1+\cdots+d_b p_b}{t} \\ \text{s.t.} \quad & p_1,\ldots,p_b\geq 0\\ & p_1+\cdots+p_b=1\\ & t\geq 0 \\ & t\leq \frac{\left( \sum_{s=1}^i p_s \right)^2}{\sum_{s=1}^i p_s L^1_{i, \{s,\dots,b\}}}, \quad i \in [b]. \end{align}\tag{36}\] This program can be written equivalently as \[\label{eq:fl2l0l1} \begin{align} \min_{q} \quad & {d_1 q_1+\cdots+d_b q_b} \\ \text{s.t.} \quad & q_1,\ldots,q_b\geq 0\\ & g_i(q) :=\left( \sum_{s=1}^i q_s \right)^2 - \sum_{s=1}^i q_s L^1_{i, \{s,\dots,b\}} \geq 0, \quad i \in [b]. \end{align}\tag{37}\] We now write the KKT conditions for 37 . Introduce multipliers \(\lambda_i\ge0\) for the constraints \(g_i(q)\ge0\) and multipliers \(\eta_k\ge0\) for the non-negativity constraints \(q_k\ge0\). The Lagrangian is \[\begin{align} \mathcal{L}(q,\lambda,s) = \sum_{i=1}^b d_i q_i - \sum_{i=1}^b \lambda_i g_i(q) - \sum_{k=1}^b \eta_k q_k. \end{align}\] Fix \(k\in[b]\). Differentiating \(g_i\) with respect to \(q_k\) yields \[\label{eq:dgi95dqk} \frac{\partial g_i(q)}{\partial q_k} = \begin{cases} 2 \sum_{s=1}^i q_s - L^1_{i,\{k,\dots,b\}}, & k\leq i,\\[4pt] 0, & k>i. \end{cases}\tag{38}\] Thus, stationarity \(\nabla_q \mathcal{L}(q,\lambda,s)=0\) yields, component-wise for \(k\in[b]\), \[\begin{align} \label{eq:stationarity} 0 = d_k - \sum_{i=k}^b \lambda_i \left( 2 \sum_{s=1}^i q_s - L^1_{i,\{k,\dots,b\}} \right) - \eta_k, \end{align}\tag{39}\] and complementarity and sign conditions are \[\begin{align} \lambda_i g_i(q) = 0, \qquad \eta_i q_i = 0, \qquad \lambda_i, \eta_i \geq 0 \qquad i\in[b]. \end{align}\]

11.2.2 \(\tau\)-Nice Sampling↩︎

For every \(i\in[b]\) we have \[\begin{align} {\mathbb{P}}\left(i \in \hat{S}\right)=\frac{\binom{b-1}{\tau-1}}{\binom{b}{\tau}}=\frac{\tau}{b} \end{align}\] and \[\begin{align} {\mathbb{E}}\left[L^{\alpha}_{i,\hat{S}} \mathbb{I}\left(i\in\hat{S}\right)\right] = \frac{1}{\binom{b}{\tau}} \sum_{\substack{S\subseteq[b]\\|S|=\tau}} L^{\alpha}_{i,S} \mathbb{I}\left(i\in S\right). \end{align}\] for \(\alpha\in\{0,1\}\). Recall that \[\begin{align} {\mathbb{E}}\left[\mathrm{cost}(S)\right] = c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( 1 - \frac{\binom{b-j}{\tau}}{\binom{b}{\tau}} \right) + \frac{\tau}{b} \sum_{j=1}^b c_j^{\sharp}. \end{align}\] Hence, substituting into 34 and 35 , the objective functions to minimize are \[\begin{align} &&{\mathbb{E}}\left[\mathrm{cost}_{\varepsilon^2}(K)\right] \\ &\propto& \frac{\sum_{i=1}^b \frac{{\mathbb{P}}\left(i \in \hat{S}\right)^2 {\mathbb{E}}\left[L^0_{i,\hat{S}} \mathbb{I}\left(i \in \hat{S}\right)\right]}{\left( {\mathbb{E}}\left[L^1_{i,\hat{S}} \mathbb{I}\left(i \in \hat{S}\right)\right] \right)^2}}{\min_{i\in[b]} \left( \frac{{\mathbb{P}}\left(i \in \hat{S}\right)^2}{{\mathbb{E}}\left[L^1_{i,\hat{S}} \mathbb{I}\left(i \in \hat{S}\right)\right]} \right)^2} \left( c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( 1 - \frac{\binom{b-j}{\tau}}{\binom{b}{\tau}} \right) + \frac{\tau}{b} \sum_{j=1}^b c_j^{\sharp} \right) \\ &=& \frac{\max_{i\in[b]} \left( \sum_{\substack{S\subseteq[b]\\|S|=\tau}} L^1_{i,S} \mathbb{I}\left(i\in S\right) \right)^2}{\frac{\tau}{b} \binom{b-1}{\tau-1}} \\ &&\qquad\times\sum_{i=1}^b \frac{\sum_{\substack{S\subseteq[b]\\|S|=\tau}} L^0_{i,S} \mathbb{I}\left(i\in S\right)}{\left( \sum_{\substack{S\subseteq[b]\\|S|=\tau}} L^1_{i,S} \mathbb{I}\left(i\in S\right) \right)^2} \left( c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( 1 - \frac{\binom{b-j}{\tau}}{\binom{b}{\tau}} \right) + \frac{\tau}{b} \sum_{j=1}^b c_j^{\sharp} \right) \end{align}\] and \[\begin{align} {\mathbb{E}}\left[\mathrm{cost}_{\varepsilon}(K)\right] &\propto& \frac{1}{\min_{i\in[b]} \frac{{\mathbb{P}}\left(i \in \hat{S}\right)^2}{{\mathbb{E}}\left[L^1_{i,\hat{S}} \mathbb{I}\left(i \in \hat{S}\right)\right]}} \left( c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( 1 - \frac{\binom{b-j}{\tau}}{\binom{b}{\tau}} \right) + \frac{\tau}{b} \sum_{j=1}^b c_j^{\sharp} \right) \\ &=& \frac{\max_{i\in[b]} \sum_{\substack{S\subseteq[b]\\|S|=\tau}} L^1_{i,S} \mathbb{I}\left(i\in S\right)}{\left( \frac{\tau}{b} \right)^2 \binom{b}{\tau}} \left( c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( 1 - \frac{\binom{b-j}{\tau}}{\binom{b}{\tau}} \right) + \frac{\tau}{b} \sum_{j=1}^b c_j^{\sharp} \right). \end{align}\] As in 11.1.2, in general there is no closed-form expression for \(\tau\) minimizing these costs, but it can be shown that \(\tau=b\) need not be optimal. Let us make a simplifying assumption that \(L^0_{i,S}\) depends only on \(i\) and \(|S|\) and define \(L^0_{i,\tau} :=L^0_{i,S}\) for \(|S|=\tau\). Then \[\begin{align} &&{\mathbb{E}}\left[\mathrm{cost}_{\varepsilon^2}(K)\right] \\ &\propto& \frac{\max_{i\in[b]} \left( \sum_{\substack{S\subseteq[b]\\|S|=\tau}} L^1_{i,\tau} \mathbb{I}\left(i\in S\right) \right)^2}{\frac{\tau}{b} \binom{b-1}{\tau-1}} \\ &&\qquad\times\sum_{i=1}^b \frac{\sum_{\substack{S\subseteq[b]\\|S|=\tau}} L^0_{i,\tau} \mathbb{I}\left(i\in S\right)}{\left( \sum_{\substack{S\subseteq[b]\\|S|=\tau}} L^1_{i,\tau} \mathbb{I}\left(i\in S\right) \right)^2} \left( c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( 1 - \frac{\binom{b-j}{\tau}}{\binom{b}{\tau}} \right) + \frac{\tau}{b} \sum_{j=1}^b c_j^{\sharp} \right) \\ &=& \frac{\max_{i\in[b]} \left( L^1_{i,\tau} \binom{b-1}{\tau-1} \right)^2}{\frac{\tau}{b} \binom{b-1}{\tau-1}} \sum_{i=1}^b \frac{L^0_{i,\tau} \binom{b-1}{\tau-1}}{\left( L^1_{i,\tau} \binom{b-1}{\tau-1} \right)^2} \left( c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( 1 - \frac{\binom{b-j}{\tau}}{\binom{b}{\tau}} \right) + \frac{\tau}{b} \sum_{j=1}^b c_j^{\sharp} \right) \\ &=& \underbrace{\max_{i\in[b]} \left( L^1_{i,\tau} \right)^2 \sum_{i=1}^b \frac{L^0_{i,\tau}}{\left( L^1_{i,\tau} \right)^2}}_{:=A_{\varepsilon^2}(\tau)} \underbrace{\left( \frac{b}{\tau} c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( \frac{b}{\tau} - \frac{\binom{b-j}{\tau}}{\binom{b-1}{\tau-1}} \right) + \sum_{j=1}^b c_j^{\sharp} \right)}_{:=B(\tau)} \end{align}\] and \[\begin{align} {\mathbb{E}}\left[\mathrm{cost}_{\varepsilon}(K)\right] &\propto& \frac{\max_{i\in[b]} \sum_{\substack{S\subseteq[b]\\|S|=\tau}} L^1_{i,\tau} \mathbb{I}\left(i\in S\right)}{\left( \frac{\tau}{b} \right)^2 \binom{b}{\tau}} \left( c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( 1 - \frac{\binom{b-j}{\tau}}{\binom{b}{\tau}} \right) + \frac{\tau}{b} \sum_{j=1}^b c_j^{\sharp} \right) \\ &=& \frac{\max_{i\in[b]} L^1_{i,\tau} \binom{b-1}{\tau-1}}{\left( \frac{\tau}{b} \right)^2 \binom{b}{\tau}} \left( c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( 1 - \frac{\binom{b-j}{\tau}}{\binom{b}{\tau}} \right) + \frac{\tau}{b} \sum_{j=1}^b c_j^{\sharp} \right) \\ &=& \underbrace{\max_{i\in[b]} L^1_{i,\tau}}_{:=A_{\varepsilon}(\tau)} \underbrace{\left( \frac{b}{\tau} c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( \frac{b}{\tau} - \frac{\binom{b-j}{\tau}}{\binom{b-1}{\tau-1}} \right) + \sum_{j=1}^b c_j^{\sharp} \right)}_{:=B(\tau)}. \end{align}\] We have already established in 11.1.2 that \(B\) is decreasing in \(\tau\). In fact, the expression to be minimized there is structurally very similar to these obtained above. Specifically, in the smooth case we had \(\mathrm{cost}_{\varepsilon}(\mathcal{D}) \propto A(\tau) B(\tau)\), where \(A(\tau) :=\max_{i\in[b]} L^0_{i,\tau}\). Since \(L^1_{i,\tau}\), like \(L^0_{i,\tau}\), is non-decreasing in \(\tau\), the same reasoning as in 11.1.2 applies to \({\mathbb{E}}\left[\mathrm{cost}_{\varepsilon}(K)\right]\). On the other hand, the dependence of \(A_{\varepsilon^2}(\tau)\) on \(\tau\) can be arbitrary depending on the scaling between \(L^0_{i,\tau}\) and \(L^1_{i,\tau}\).

Consequently, both objectives \({\mathbb{E}}\left[\mathrm{cost}_{\varepsilon}(K)\right]\) and \({\mathbb{E}}\left[\mathrm{cost}_{\varepsilon^2}(K)\right]\) present a trade-off between a decreasing factor \(B(\tau)\) and a (typically) increasing factor \(A_1(\tau)\). If \(A_{\varepsilon}(\tau)\) or \(A_{\varepsilon^2}(\tau)\) grows sufficiently fast in \(\tau\), it can compensate for the decrease of \(B(\tau)\) and make a smaller \(\tau\) (or even \(\tau=1\)) optimal.

11.2.3 \(\tau\)-Submodel Sampling↩︎

For any fixed \(i\in[b]\), we have \[\begin{align} {\mathbb{P}}\left(i \in \hat{S}\right) &=& \sum_{j=\max\{1,i-\tau+1\}}^{\min\{i,b-\tau+1\}} p_j \end{align}\] and \[\begin{align} {\mathbb{E}}\left[L^1_{i,\hat{S}} \mathbb{I}\left(i \in \hat{S}\right)\right] &=& \sum_{j=\max\{1,i-\tau+1\}}^{\min\{i,b-\tau+1\}} p_j L^1_{i,\{j,\dots,j+\tau-1\}}. \end{align}\] We have also already shown that \[\begin{align} {\mathbb{E}}\left[\mathrm{cost}(\hat{S})\right] = c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( \sum_{i=1}^{\min\{j,b-\tau+1\}} p_i \right) + \sum_{j=1}^b c_j^{\sharp} \left( \sum_{i=\max\{1,j-\tau+1\}}^{\min\{j,b-\tau+1\}} p_i \right). \end{align}\]

11.2.3.1 Partitioned \(\tau\)-submodel sampling.

Mimicking 11.1.3.1, let us consider a special case of the above sampling scheme where the submodels partition \([b]\). For simplicity, we assume that \(b\) is divisible by \(\tau\) and let \(m=\frac{b}{\tau}\). The algorithm is then equivalent to \(\tau\)-submodel sampling with starting-index distribution satisfying \[\begin{align} p_{s_k} > 0 \quad (k=1,\dots,m), \qquad p_j=0 \text{ otherwise}, \end{align}\] where \(s_k=(k-1)\tau+1\), \(k\in[m]\). For any fixed \(i\in[b]\), the derivations above simplify to \[\begin{align} {\mathbb{P}}\left(i \in \hat{S}\right) = \sum_{j=\max\{1,i-\tau+1\}}^{\min\{i,b-\tau+1\}} p_j = p_{s_{\lceil i/\tau \rceil}} \end{align}\] and \[\begin{align} {\mathbb{E}}\left[L^{\alpha}_{i,\hat{S}} \mathbb{I}\left(i \in \hat{S}\right)\right] = \sum_{j=\max\{1,i-\tau+1\}}^{\min\{i,b-\tau+1\}} p_j L^{\alpha}_{i,\{j,\dots,j+\tau-1\}} = p_{s_{\lceil i/\tau \rceil}} L^{\alpha}_{i,B_{\left\lceil \frac{i}{\tau} \right\rceil}} \end{align}\] for \(\alpha\in\{0,1\}\). Plugging this choice into 34 and 35 gives \[\begin{align} {\mathbb{E}}\left[\mathrm{cost}_{\varepsilon^2}(K)\right] &\propto& \frac{\sum_{i=1}^b \frac{{\mathbb{P}}\left(i\in \hat{S}\right)^2 {\mathbb{E}}\left[L^0_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right]}{\left( {\mathbb{E}}\left[L^1_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right] \right)^2}}{\min_{i\in[b]} \left( \frac{{\mathbb{P}}\left(i\in \hat{S}\right)^2}{{\mathbb{E}}\left[L^1_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right]} \right)^2} \\ &&\qquad\times\left( c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( \sum_{i=1}^{\min\{j,b-\tau+1\}} p_i \right) + \sum_{j=1}^b c_j^{\sharp} \left( \sum_{i=\max\{1,j-\tau+1\}}^{\min\{j,b-\tau+1\}} p_i \right) \right) \\ &=& \frac{\sum_{i=1}^b \frac{p_{s_{\lceil i/\tau \rceil}} L^0_{i,B_{\left\lceil \frac{i}{\tau} \right\rceil}}}{\left( L^1_{i,B_{\left\lceil \frac{i}{\tau} \right\rceil}} \right)^2}}{\min_{i\in[b]} \left( \frac{p_{s_{\lceil i/\tau \rceil}}}{L^1_{i,B_{\left\lceil \frac{i}{\tau} \right\rceil}}} \right)^2} \left( c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( \sum_{k=1}^{\left\lceil \frac{j}{\tau} \right\rceil} p_{s_k} \right) + \sum_{j=1}^b c_j^{\sharp} p_{s_{\left\lceil \frac{j}{\tau} \right\rceil}} \right) \end{align}\] and \[\begin{align} {\mathbb{E}}\left[\mathrm{cost}_{\varepsilon}(K)\right] &\propto& \frac{1}{\min_{i\in[b]} \frac{{\mathbb{P}}\left(i\in \hat{S}\right)^2}{{\mathbb{E}}\left[L^1_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right]}} \\ &&\qquad\times\left( c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( \sum_{i=1}^{\min\{j,b-\tau+1\}} p_i \right) + \sum_{j=1}^b c_j^{\sharp} \left( \sum_{i=\max\{1,j-\tau+1\}}^{\min\{j,b-\tau+1\}} p_i \right) \right) \\ &=& \frac{1}{\min_{i\in[b]} \frac{p_{s_{\lceil i/\tau \rceil}}}{L^1_{i,B_{\left\lceil \frac{i}{\tau} \right\rceil}}}} \left( c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( \sum_{k=1}^{\left\lceil \frac{j}{\tau} \right\rceil} p_{s_k} \right) + \sum_{j=1}^b c_j^{\sharp} p_{s_{\left\lceil \frac{j}{\tau} \right\rceil}} \right). \end{align}\] The expression for \({\mathbb{E}}\left[\mathrm{cost}_{\varepsilon}(K)\right]\) is entirely analogous to 24 derived in the smooth case, with \(L^0_{i,B_{\left\lceil \frac{i}{\tau} \right\rceil}}\) replaced by \(L^1_{i,B_{\left\lceil \frac{i}{\tau} \right\rceil}}\). Consequently, for a fixed \(\tau\), the optimal block probabilities are \[\begin{align} p_{s_k}^\star = \frac{\max_{i\in B_k} L^1_{i,B_k}}{\sum_{l=1}^m \max_{i\in B_l} L^1_{i,B_l}}, \end{align}\] Again, the optimal number of blocks \(m\) depends on the growth rate of the constants \(L^1_{i,\tau}\) with respect to \(\tau\) and on their interaction with the costs \(c_j\) and \(c_j^{\sharp}\). In general, the best choice of \(m\) is not necessarily \(b\).

The dependence of \({\mathbb{E}}\left[\mathrm{cost}_{\varepsilon^2}(K)\right]\) on \(\tau\) is significantly more complex and governed by the relative scaling between \(L^0_{i,\tau}\) and \(L^1_{i,\tau}\).

11.2.4 Arbitrary Submodel Sampling↩︎

For any \(i\in[b]\), we have \[\begin{align} {\mathbb{P}}\left(i \in \hat{S}\right) = p_{k(i)} \end{align}\] and \[\begin{align} {\mathbb{E}}\left[L^{\alpha}_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right] = p_{k(i)} L^{\alpha}_{i,B_{k(i)}} \end{align}\] for \(\alpha\in\{0,1\}\). Since \[\begin{align} {\mathbb{E}}\left[\mathrm{cost}(\hat{S})\right] = c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( \sum_{k:\underline{b}_k \leq j} p_k \right) + \sum_{j=1}^b c_j^{\sharp} p_{k(j)}, \end{align}\] the total expected costs become \[\begin{align} {\mathbb{E}}\left[\mathrm{cost}_{\varepsilon^2}(K)\right] &\propto& \frac{\sum_{i=1}^b \frac{{\mathbb{P}}\left(i\in \hat{S}\right)^2 {\mathbb{E}}\left[L^0_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right]}{\left( {\mathbb{E}}\left[L^1_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right] \right)^2}}{\min_{i\in[b]} \left( \frac{{\mathbb{P}}\left(i\in \hat{S}\right)^2}{{\mathbb{E}}\left[L^1_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right]} \right)^2} \left( c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( \sum_{k:\underline{b}_k \leq j} p_k \right) + \sum_{j=1}^b c_j^{\sharp} p_{k(j)} \right) \\ &=& \frac{\sum_{i=1}^b \frac{p_{k(i)} L^0_{i,B_{k(i)}}}{\left( L^1_{i,B_{k(i)}} \right)^2}}{\min_{i\in[b]} \left( \frac{p_{k(i)}}{L^1_{i,B_{k(i)}}} \right)^2} \left( c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( \sum_{k:\underline{b}_k \leq j} p_k \right) + \sum_{j=1}^b c_j^{\sharp} p_{k(j)} \right) \end{align}\] and \[\begin{align} {\mathbb{E}}\left[\mathrm{cost}_{\varepsilon}(K)\right] &\propto& \frac{1}{\min_{i\in[b]} \frac{{\mathbb{P}}\left(i\in \hat{S}\right)^2}{{\mathbb{E}}\left[L^1_{i,\hat{S}} \mathbb{I}\left(i\in \hat{S}\right)\right]}} \left( c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( \sum_{k:\underline{b}_k \leq j} p_k \right) + \sum_{j=1}^b c_j^{\sharp} p_{k(j)} \right) \\ &=& \frac{1}{\min_{i\in[b]} \frac{p_{k(i)}}{L^1_{i,B_{k(i)}}}} \left( c_{\mathrm{ov}} + \sum_{j=1}^b c_j \left( \sum_{k:\underline{b}_k \leq j} p_k \right) + \sum_{j=1}^b c_j^{\sharp} p_{k(j)} \right). \end{align}\] Following the reasoning in 11.1.4, we find that for a fixed partition, the choice of probabilities minimizing \({\mathbb{E}}\left[\mathrm{cost}_{\varepsilon}(K)\right]\) is \[\begin{align} p_k^\star = \frac{\frac{1}{\delta_k^{\min}}}{\sum_{l=1}^m \frac{1}{\delta_l^{\min}}} = \frac{\max_{i\in B_k} L^0_{i,B_k}}{\sum_{l=1}^m \max_{i\in B_l} L^0_{i,B_l}}, \end{align}\] and that submodel training can improve upon full model training in certain regimes.

12 Convergence Results – Stochastic Gradient Case↩︎

The proof of 16 relies on two preliminary lemmas, which we state and prove first.

13 Experiments↩︎

We evaluate Drop-Muon on three standard benchmarks–MNIST, Fashion-MNIST, and CIFAR-10–using 3-layer convolutional neural networks (CNNs) of varying capacity. Our goal is to study how partial layer updates of Drop-Muon accelerate training relative to standard Muon (a special case of Drop-Muon with full-network updates).

Experiments were run on various NVIDIA GPUs based on availability. MNIST and Fashion-MNIST experiments were run on Tesla V100-SXM2-32GB. CIFAR-10 experiments were performed on a mix of A100-SXM4-80GB, Tesla P100-PCIE-16GB, and GeForce GTX 1080 Ti GPUs.

13.1 Key Training Parameters↩︎

We identify several key training parameters and conduct a grid search across these parameters to systematically explore their effects.

13.1.0.1 Layer sampling distribution.

The choice of layer distribution determines which layers are updated at each iteration and is central to the behavior of Drop-Muon. We considered the following strategies:

• Uniform Distribution: Each layer is sampled with equal probability (\(p_i = \frac{1}{b}\) for all \(i\in[b]\) in the notation from 4).

• Linear Distribution: Sampling probability increases linearly with the layer index, favoring deeper layers. Empirically, this distribution performed poorly.

• Quadratic Distribution: Probability increases quadratically with depth, strongly biasing toward deeper layers. Similarly, this performed poorly.

• Exponential Distribution: Layer probability grows exponentially with depth, emphasizing deeper layers. This also proved suboptimal.

• Epoch-Shift Distribution: The distribution biases the sampling towards shallow layers at the beginning of training, gradually shifting towards deeper layers as the training progresses (see more details in 13.2). This dynamic adjustment might help in balancing early feature extraction and later complex feature learning during the optimization process–see example in 6.

13.1.0.2 Batch size and learning rate.

We test batch sizes \(\{64, 512, 8192, 16384, 32768\}\) and learning rates \(\{0.1, 0.01, 0.001\}\).

13.1.0.3 Model depth and capacity.

To assess the impact of network complexity, we train multiple 3-layer CNNs with varying channel configurations: \([8,16,32]\), \([16,32,64]\), \([64,128,256]\), \([128,256,512]\), \([64,16,8]\), and \([256,128,64]\).

13.1.0.4 Fixed parameters.

We fix the number of epochs to \(20\) or \(50\), momentum to \(0.5\), and the number of Newton-Schulz iterations to \(5\) (see 7.1).

This grid search allows us to systematically evaluate Drop-Muon. While many combinations performed poorly (and are therefore omitted in the results that follow), the experiments highlighted configurations–particularly with uniform or epoch-shift distributions–that consistently reduce training time while maintaining final accuracy.

13.2 Epoch-Shift Distribution↩︎

Drop-Muon can dynamically adjust which layers are updated at each iteration. The strategy we consider is the epoch-shift distribution, which biases training towards shallow layers in the early epochs, gradually shifting focus to deeper layers as training progresses (see 6). Such a schedule balances early feature extraction with later complex feature learning, improving convergence in practice.

Formally, let \(b\) denote the total number of layers, \(\alpha\) be a constant controlling the sharpness of the bias, and define the training progress as \[\text{progress} = \frac{\text{epoch}}{\text{max\_epochs}} \in [0,1].\] The weight for each layer \(i\) is given by \[w_i = \exp \left( \alpha \left[ (1 - \text{progress})(b - 1 - i) + \text{progress} \cdot i \right] \right).\] We then normalize the weights to obtain a valid probability distribution: \[W = \sum_{i=0}^{b-1} w_i, \qquad p_i = \frac{w_i}{W}.\]

By adjusting \(\alpha\), one can control the rate at which training emphasis shifts from shallow to deeper layers. In 6, we set \(\alpha = 0.5\).

13.3 Evaluation Metrics↩︎

We repeat every experiment over multiple random seeds. When reporting aggregated results, we use two complementary procedures:

• Normalized curve averaging: We normalize wall-clock time for each run so that Muon always ends at \(t=1\), interpolate both Muon and Drop-Muon curves onto this grid, and average over seeds.

• Time-to-target evaluation: We measure the wall-clock time to reach fixed accuracy thresholds (e.g. \(60\%\),\(70\%\),\(80\%\),...), reporting the ratio between Muon and Drop-Muon. This metric is interpolation-free and directly quantifies practical speedup.

13.4 Results on MNIST↩︎

We first evaluate Drop-Muon on the MNIST dataset. Figure 2 shows a representative run comparing standard full-network Muon training with Drop-Muon using a uniform layer sampling distribution. Although the per-epoch train accuracy of Drop-Muon is initially lower than that of full-layer Muon, the wall-clock training time tells a different story: for training times up to approximately 150 seconds, Drop-Muon consistently achieves higher accuracy. In practice, this means that reaching a train accuracy of \(95\%\) or less is faster with Drop-Muon, highlighting its efficiency advantage.

To account for variability across independent runs, we evaluate two aggregation strategies. 7 shows normalized curve averaging, where each run is rescaled to a common time grid. Figure 3 (left) presents averaged time-to-target evaluation across multiple seeds, reporting the ratio of Muon time to Drop-Muon time for fixed accuracy thresholds.

The time-to-target results clearly demonstrate a speedup of up to \(1.4\times\) across thresholds of \(60\%, 70\%, 80\%\), and \(99\%\), with slightly lower speedup for the \(90\%\) threshold. We note variability across seeds, emphasizing the stochastic nature of layer sampling. Despite this, Drop-Muon consistently provides practical acceleration across training.

13.5 Results on Fashion-MNIST↩︎

Figure 8 shows a typical run comparing standard Muon training with Drop-Muon using the epoch-shift layer sampling distribution. Early in training, Drop-Muon achieves faster progress in wall-clock time, even if per-epoch accuracy is slightly lower.

Normalized curve averaging across multiple seeds is presented in 9. Although the curves appear smoother at higher accuracy, this should not be interpreted as a reduction in variance; early-stage training is inherently noisier due to layer sampling, whereas later training reflects fewer updates per wall-clock time.

Figure 3 (right) reports the averaged time-to-target evaluation. Drop-Muon consistently achieves speedups of approximately \(1.2\) across all considered accuracy thresholds. The \(99\%\) threshold is omitted as neither Muon nor Drop-Muon reach this level within the plotted runs.

Overall, Fashion-MNIST experiments reinforce that Drop-Muon provides consistent practical speedups, similar to MNIST, while maintaining comparable final accuracy.

13.6 Results on CIFAR-10↩︎

CIFAR-10 requires more sophisticated network architectures. We use a CNN with channels \([128,256,512]\) and insert batch normalization layers between each convolution and ReLU activation to improve training stability.

4 shows an example run comparing standard full-layer Muon training with Drop-Muon using the epoch-shift layer sampling distribution. While Drop-Muon has slightly lower per-epoch training accuracy, it achieves faster progress in terms of wall-clock time. 10 shows normalized curve averaging across multiple seeds (the flat end of the curve reflects the earlier completion of Drop-Muon).

Time-to-target results in 11 indicate a notable speedup at the \(90\%\) train accuracy threshold. However, for lower thresholds (\(60\%, 70\%, 80\%\)), speedups are minimal or absent for this configuration. We attribute this limitation primarily to sub-optimal hyperparameter choices (see 6).

13.7 Discussion and Practical Remarks↩︎

We summarize several key observations and practical considerations arising from our experiments:

• Simplicity of implementation: Drop-Muon can be implemented in just a few lines of code, making it easy to integrate into existing training pipelines.

• Further benefits from tuning: Performance can be further improved through dedicated tuning of hyperparameters (see 6).

• Potential for implementation improvements:

  • Turning off gradient computations at every iteration can be costly for large models. Sampling the cutoff layer only every few iterations can reduce overhead and improve efficiency.

  • When the batch size or model size is small, the gradient computation may be dominated by the Newton-Schulz routine. In such cases, computing all gradients but orthogonalizing only a subset can be advantageous.

• Variance across seeds: The method exhibits non-negligible variance, particularly on smaller datasets. Stabilization mechanisms could mitigate this.

• Adaptive learning of sampling distributions: A promising future direction is to learn the layer sampling distribution online, so that it automatically adapts to the training dynamics of a given dataset and architecture.

14 Useful Facts↩︎

For all \(X,Y \in \mathcal{X}\) and \(t>0\), we have: \[\begin{align} \tag{47} \left\| {\rm LMO}_{\mathcal{B}(0,t)}(G) \right\| &= t \\ \tag{48} \left\langle G,{\rm LMO}_{\mathcal{B}(X,t)}(G)\right\rangle &= -t \left\| G \right\|_\star \\ \tag{49} \left\langle X,X^\sharp\right\rangle &= \left\| X^\sharp \right\|^2, \\ \tag{50} \left\| X \right\|_{\star} &= \left\| X^\sharp \right\|, \\ \tag{51} \left\langle H,X\right\rangle &= \left\| X \right\|_{\star}, \quad \left\| H \right\| = 1, \qquad\forall X\neq0. \end{align}\] where \(H \in \partial\left\| \cdot \right\|_{\star}(X)\) belongs to the subdifferential of the dual norm.

References↩︎