ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
Authors: Arto Maranjyan
El Mehdi Saad
Peter Richtárik
Francesco Orabona
Venue: ICML 2025 poster (2025)
Keywords: Multi-Armed Bandit
UCB
adaptive task allocation
asynchronous methods
parallel methods
SGD
Primary area: theory->online_learning_and_bandits
TL;DR: We propose a method that learns machine speeds on the fly to assign tasks more efficiently in parallel computing.

Abstract: Asynchronous methods are fundamental for parallelizing computations in distributed machine learning. 
    They aim to accelerate training by fully utilizing all available resources.
    However, their greedy approach can lead to inefficiencies using more computation than required, especially when computation times vary across devices.
    If the computation times were known in advance, training could be fast and resource-efficient by assigning more tasks to faster workers.
    The challenge lies in achieving this optimal allocation without prior knowledge of the computation time distributions.
    In this paper, we propose ATA (Adaptive Task Allocation), a method that adapts to heterogeneous and random distributions of worker computation times.
    Through rigorous theoretical analysis, we show that ATA identifies the optimal task allocation and performs comparably to methods with prior knowledge of computation times.
    Experimental results further demonstrate that ATA is resource-efficient, significantly reducing costs compared to the greedy approach, which can be arbitrarily expensive depending on the number of workers.

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
Artavazd Maranjyan'
El Mehdi Saad
Peter Richtárik Francesco Orabona
Abstract
In order to perform single iteration of the algorithm, a certain number (B) of tasks needs to be performed.
Asynchronous methods are fundamental for par- allelizing computations in distributed machine learning. They aim to accelerate training by fully utilizing all available resources. However, their greedy approach can lead to inefficiencies using more computation than required, especially when computation times vary across devices. If the computation times were known in advance, train- ing could be fast and resource-efficient by assign- ing more tasks to faster workers. The challenge lies in achieving this optimal allocation without prior knowledge of the computation time distri- butions. this paper, we propose ATA (Adaptive Task Allocation), method that adapts het- erogeneous and random distributions of worker computation times. Through rigorous theoretical analysis, we show that ATA identifies the optimal task allocation and performs comparably to meth- ods with prior knowledge of computation times. Experimental results further demonstrate that ATA is resource- -esfurceefficceet, significantly reducing costs compared to the greedy approach, which be arbitrarily expensive depending the number of workers.
Each task can be computed by any worker, and the tasks are not temporally related. That is, they can be computed in any order, in parallel, and so on.
Whenever a worker is asked to perform single task, the task will take certain amount of time, modeled as a nonnegative random variable drawn from an un- known distribution specific to that worker. The stochas- tic assumption makes sense because in real systems computation times are not fixed and can vary with each iteration (Dean & Barroso, 2013; Chen et al., 20; Dutta et al., 2018; Maranjyan et al., 2025a).
Each worker can only work single task at time. That is, worker processes all tasks it has to perform sequentially. Different workers work in parallel.
A natural goal in this setup is to make sure all tasks are completed as fast as possible (in expectation), which mini- mizes the (expected) time takes for single iteration of the algorithm to be performed provided that the task completion time is the dominant time factor of the iteration. Provided are willing to waste resources, there is simple solution to this problem, a Greedy Task Allocation (GTA) strategy, which follows this principle: Make sure all workers are always busy working on some task, and stop once B tasks have been completed. GTA, we initially ask all n workers to start working on task, and as soon as some worker is done with task, ask it to start completing another task. This process is repeated until B tasks have been completed.
1. Introduction
In this work, we address a very general yet fundamental and important problem arising in various contexts and fields. In particular, there workers/nodes/deviees collaborating to run some iterative algorithm which has the following structure:
While GTA minimizes the completion time, it can be im- mensely wasteful in terms the total worker utilization time needed to collect all B tasks. Indeed, consider the scenario with n = 1000 workers and = 10 tasks. this case, GTA will lead to at least n B 990 unnecessary tasks being run in each iteration! This is highly undesirable in situations where the workers are utilized across multiple other jobs besides running the iterative algorithm mentioned above.
'King Abdullah University Science and Technol- ogy (KAUST), Thuwal, Saudi Arabia. Correspondence to: Artavazd Maranjyan Sarto.maranjyanee @ @aaaal.mommm.......... El Mehdi Saad <mehdi.saad @kasst..duuss Peter Richtárik <peter.richtarik @kwwst.cdu.sa.. <francesco www.eera.bbbaa.........
Francesco Orabona
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
The goal of our work is to design new task allocation strate- gies, with rigorous theoretical support, that would attempt

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
to minimize the expected completion time subject to the constraint that such wastefulness is completely eliminated. That is, we ensure that more than B tasks are completed in each round.
the proposed allocation problem. ATA is agnostic to work- ers' computation times, and our theoretical analysis demon- strates that the total computation time achieved by our meth- ods remains within a small multiplicative factor of the op- timal computation time (i.e., the one attainable with full knowledge of the workers' arm distributions). Additionally, we present ATA- ATA-iiccii,,,, a variant of ATA that leverages a novel data-dependent concentration inequality and achieves better empirical results. Finally, we validate our approach through numerical simulations.
1.1. A Motivating Example: Optimal Parallel SGD
A key inspiration for our work, and the prime example of the general task collection problem described above, relates to recent development in the area of parallel stochastic gradient descent (SGD) methods. Consider the problem of finding an approximate stationary point of the optimization problem
2. Related Work
min (f(x)= := E [Fe( #ER"
Most of the literature on asynchronous methods focuses on demonstrating advantages over their synchronous counter- parts. For the simplest method, SGD, this was only recently established by Tyurin & Richtárik (2024). With this result in place, the community can now shift its focus to reducing the overhead of asynchrony. Our work may be the first step in this direction.
where Se: IR + R are smooth nonconvex functions, and I is assumed to be bounded from below. We assume that
EEED [I]e(()) - [[(=))]]] s²²
for all x E Rd.
In a recent breakthrough, Tyurin & Richtárik (2024) re- cently developed a parallel SGD method, optimal in terms of novel notion of complexity called time complexity, for solving the above problem with n parallel workers, assum- ing that it takes т; >> O seconds to worker to compute a stochastic gradient of f (this corresponds to a task). Their method, Rennala SGD, corresponds to Minibatch SGD of minibatch size B (which depends on the target accuracy and only), with the B tasks (stochastic gradients) completed via GTA. While minimax optimal in terms of time com- plexity, the GTA task allocation strategy employed within Rennala SGD can be wasteful, as explained above.
In federated learning (FL) (Konečný et al., 2016; McMahan et al., 2016; Kairouz et al., 2021), several works account for system heterogeneity. The most well-known FL method, FedAvg (McMahan et al., 2017), operates by performing multiple local steps on workers, where each step can be viewed as task. Some works adjust the number of local steps based on worker computation times (Li et al., 2020; Maranjyan et al., 2022), effectively adapting task assign- ments to worker speed. However, these methods rely on prior knowledge of these times rather than learning them adaptively, as we do.
We reformulate our problem as an online bandit problem. The literature on bandit algorithms is vast, and we refer the reader to Lattimore & Szepesvári (2020) for an introduction to this subject. Our algorithm based on the approach of us- ing Lower Confidence Bounds (LCBs) on the true means of the arms. This idea, originally proposed by Auer (2002) for the classical Multi-Armed Bandit (MAB) setting, has since been widely adopted in the stochastic combinatorial bandits literature (Gai et al., 2012; Chen et al., 2013; Combes et al., 2015; Kveton et al., 2015). Using LCBs instead of the em- pirical estimates of the means allows to trade-off optimally exploration and exploitation.
Recently, Maranjyan et al. (2025b) proposed Ringmaster ASGD, fully asynchronous SGD method, matching the optimal time complexity of Rennala SGD and achieving op- timality for arbitrary compute time patterns associated with the tasks (stochastic gradients), including random, as consid- ered in our setup. However, Ringmaster ASGD also employs a greedy task allocation strategy, leading to wastefulness.
Numerous other parallel/distributee methods involve the im- plementation of a task allocation strategy, including stochas- tic proximal point methods (task = evaluation of the stochas- tic prox operator), higher-order methods (task evaluation of stochastic Hessian), and beyond. So, by addressing the general task allocation problem, we aim to tame the inherent resource wastefulness of all these methods.
The "greedy" approach we employ, which involves select- ing the action that minimizes the loss function based on lower confidence bounds instead of the unknown means, is a standard technique in the literature (Chen et al., 2013; Lin et al., 2015). However, note that our larger action space and the discontinuity of our loss function necessitates a more tailored analysis. To the best of our knowledge, this is the first work addressing non-continuous loss function in a stochastic combinatorial MAB-lik framework. To over-
1.2. Contributions
In this work, we formalize the task allocation problem as a combinatorial online learning problem with partial feed- back and non- linear losses. Then, we introduce ATA, a lower-confidence hound-based algorithm designed to solve
2

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
come this challenge, we exploit the specific structures of loss function and action space to control the number of rounds where suboptimal actions are chosen. Additionally, our procedure is computationally efficient.
assumption, we recall the definition of the sub-exponential norm, also known as the Orlicz norm, for centered real- valued random variable X:
| == inf > O Elexp((X|/C) 2)
(2)
3. Problem Setup
Hence, formally make the following assumption.
Assumption 3.1. Let 2 0. For all E [n], X; is positive
In this section, we formally describe the problem setup.
random variable and ||;; pill S ar
3.1. Task Allocation Protocol
In the remainder of this paper we denote or == ||Xill for each i E [n], let a := max ag
We consider system of n workers, each responsible for computing gradients. In each round, the allocation algo- rithm has a budget of B units that must be allocated among the workers. Each unit allocation will result in one gra- dient computation. We denote by K the total number of rounds, which is assumed to be unknown to the learner. We denote by x-1) the computation time of the worker i E [n] = 21. gradient. Consequently, the computation time required for worker i perform its task of computing ai gradients in round & is given by
The considered class encompasses several other well-known classes of distributions in the literature, such as support- bounded and sub-Gause Gaussian distributions. Moreover, in- cludes exponential distributions, which are frequently used in the literature to model waiting computation times in queuing theory and resource allocation in large distributed systems (Gelenbe & Mitrani, 2010; Gross et al., 201 I; Had- jis et al., 2016; Mitliagkas et al., 2016; Dutta et al., 2018; Nguyen et al., 2022).
i.k
for round & [K] on its u-th
3.3. Objective of the Allocation Algorithm
The main objective of this work is to develop an online allocation strategy with small expected total computation time, defined as
if a 2 1, and O otherwise.
In each round k. the allocation algorithm must choose an allocation vector ag E N" such that ||a||, = B, based on the information available prior to round k. The feedback consists of di observed times for all the chosen workers. We will denote the action set by
K Ck = Σ E(C(ax))
the distributions of the arms were known in advance, the optimal allocation a' E A would be selected to minimize the expected computation time per round, E[C(-)), and this allocation would be used consistently over K rounds, lead- ing to the optimal total computation time
If
A: (a E N" ||l|| = B],
where N is the set of natural numbers, including the 0.
The objective of the allocation strategy in each round &s is to minimize the total computation time. Hence, the objective is to minimize C : A + R,, the computation time that the optimizer waits to receive B gradients using an allocation vector a E A, defined as
Cx - =KE[C(a")]
Our goal is to design a strategy that ensures the computation time Cn remains within a small multiplicative factor of the optimal time Ck. plus an additional negligible term. Specifically, we aim to satisfy
C(ax)= = max x (a)
(esupp(as)) (
u=1
CK S Cx + EK.
(3)
3.2. Modeling Assumptions
where 7 2 1 is a constant close to 1, and EK is a negligible term compared to Cj when K - 0o. This would assure us that in the limit we are a constant multiplicative factor away from the performance of the optimal allocation strategy that has full knowledge of the distributions of the computational times of the workers.
We assume that the computation time of each worker i E [n]] are i.i.d. drawn from a random variable X; following probability distribution 2/2, We denote by μ4 (M1...... Mb,) the vector of unknown means. Hence, the random variables (x(()) with E (1....,0i.b) ) are 0i,k i.i.d. samples drawn from P4.
Finding strategy solving the objective in (3) presents sev- eral technical challenges. First, the action space A is dis- crete, and the nonlinearity of the computation time functions
We assume that the distribution 2/ of the computation times to be sub-exponential random variables. To quantify this
3

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
C(-) prevents reducing our objective a convex problem. Second, the size of A is combinatorial, growing the or- der of (*++ B ). which necessitates exploiting the inherent B
we show that
l(a, EE([(a]] - (1 +1+ 4n (n)))))) u))
(5)
problem structure to develop efficient strategies. Third, be- cause the workers' computation times are stochastic, any
where 1) is defined as
solution must account for uncertainty. Finally, the online setting forces the learner to balance exploration and exploita- tion under a limited allocation budget of B units per round and partial feedback only the computation times of work- ers who receive allocations are observed. This last point naturally suggests adopting a MAB approach.
17 == max iE 14
(6)
In words, 'j provides an upper bound on the ratio between the standard deviation and the mean of the arms. Note that in the literature, lang, or Gamma distributions, where the ratio v) typically! bounded by 1.
common to consider exponential, Er-
In the next section, we show how to reduce this problem to a MAB problem and how to efficiently solve it.
The bound above will allow us to derive guarantees on the total computation time of an allocation strategy based on its guarantees for the proxy loss {(-), up to a factor of the order + 4ny In( (B). We remark that in the special case where the
4. Adaptive Task Allocation
Here, we first show how to reduce the problem in (3) to non-linear stochastic Multi-Armed Bandit (MAB) problem. Then, we propose an efficient algorithm for this formulation.
arms' distributions are deterministic (17 = 0) or the query budget is unitary (B = 1), the two targets EE [C(a)] and l exactly coincide.
4.1. Reduction to Multi-Armed Bandit and Proxy Loss
4.2. Comparison with the Combinatorial Bandits Setting
The stochastic MAB problem is fundamental framework in sequential deciisionmmaking under uncertainty. It involves a scenario where an agent must choose among a set of arms, each associated with an unknown reward distribution. The agent aims to maximize cumulative reward (or equivalently minimize the cumulative loss) over time by balancing explo- ration (gathering information about the reward distributions) and exploitation (leveraging the best-known arm). The chal- lenge lies in the trade-off between exploring suboptimal arms to refine reward estimates and exploiting the arm with the highest observed reward, given the stochastic nature of the outcomes. Using the terminology from bandit literature, here we will refer each worker as an "arm."
Our setting is closely related to the Combinatorial Multi- Armed Bandits (CMAB) framework (Cesa-Bianchi & Lu- gosi, 2012), particularly due to the combinatorial nature of the action space and the semi-bandit feedback, where the learner observes outcomes from all chosen arms. However, our formulation differs in two significant ways. First, while CMAB typically involves selecting subset of arms, re- sulting in action space with a maximum size of 2" our action space A has cardinality of ("-B-1). The ratio between these two can be extremely large, potentially grow- ing exponentially with n. Second, although most works in this domain assume a linear loss function in the arms' means, some notable exceptions address non-linear reward functions (Chen et al., 2013; Lin et al., 2015: Chen et al., 201 Wang & Chen, 2018). However, these approaches generally rely on assumptions such as smoothness, Lipschitz continuity, or higher-order differentiability of the reward function. In contrast, our loss function ((.)1)) is not contin- uous with respect to the arms' means. Finally, motivated by the practical requirements of our setting, we place a strong emphasis on computational efficiency that rules out most of the approaches based on CMAB.
However, differently from the standard MAB problem, we have a harder problem because E [C(az)) depends on the joint distribution of all the arms in the support of ag, rather than on their expectations only. This dependency potentially renders the task of relying on estimates of E [C'(a)] for a E A computationally challenging due to the combinatorial nature of the set A.
To solve this issue, our first idea is to introduce a proxy loss L: : A x B550 + R2o, defined as
E(a, p1) == max 0:/4: iE[n]
(4)
4.3. Adaptive Task Allocation Algorithm
Now, we introduce our Adaptive Task Allocation algorithm (ATA). ATA does not require prior knowledge of the hori- zon K and only relies on an upper bound satisfying
Due to the convexity of C(-), the introduced proxy-loss underestimates the expected computation time. However, in Appendix D.3 we prove that this quantity also upper bounds the expected computation time up to a constant that depends on the distribution of the arms In particular, for any a E A,
For Gamma(o, A), o/p 1/va, so the claim holds for a a21 1.
4

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
Algorithm ATA (Adaptive Task Allocation)
the chosen optimization algorithm. Refer to Appendix B for further details.
1: Input: allocation budget B, > 0
2: Initialize: empirical means /4,,1 = O, usage counts
As last step, the feedback obtained after applying the alloca- tion ag is used to update the lower confidence bounds. The complete pseudocode for ATA is provided in Algorithm 1.
K = 0, and usage times T; = 0, for all i E [n]
3: for k == l....... K do
4: Compute LCBs (st,x) for all E [v>] using (7)
5: Find allocation: ag E arg min E(a, 8:)
4.4. Upper-Bound on the Total Computation Time
6: Allocate a tasks to each worker E [n]
7: Update optimization parameters
We provide guarantees for ATA in the form of an upper bound the expected total computation time required to perform K iterations of the optimization procedure. Recall that the proxy loss E(., (r) and the expected computation time are related through (5). This relationship and Theorem 6.1 allow us to derive guarantees on the expected total computation time, denoted by
8: for i such that a o 0 do
9:
K1 = Ki + 0i,k Th,+11 = Ti + -i=
10:
11: A1,,+1 = TrA+1/ Ki,t+1
12: end for
13:
end for
K
CK B|C(ax)]
a' 2 max ||x;- - pille, for the Orlicz norms of the arm
distributions. Recall that [_X;- pill,, I 2 || p so an upper bound on ||, also provides one for |X; pallg.. The core idea of the procedure is to allocate the workers based on lower confidence bound estimates on the arm means (μ4)sE|)) in order to balance exploration and ex- ploitation.
We define the optimal allocation for minimizing the compu- tation time
mig a" E arg min E [C([)]]
aEA
Consequently, the optimal expected total computation time in this framework is given by
For each arm i E [ni] and round E [K] let Kr, represent the number of samples collected from the distribution of arm i up to round k. At each round k. compute empirical mean, denoted by jt, k. using the K samples obtained so far. Based on these empirical means, we define the lower confidence bounds 8i,k as
Ck = KE(C(a"))
Theorem 4.2 (Proof in Appendix D.3). Suppose Assump- tion 3.1 holds and let 1 == max;e|n] 0i/ Then, the total expected computation time after K rounds, using the allo- cation prescribed by ATA with inputs (B,a) satisfies
Si,k (0s,k conf(i, (i, 4))4
where (x)+ maxfar, 0} and conf(:,.) is defined as
CK S (1 +1+ 4n In(B)) + O(InK)
cons(i, [ 2a (v
[n(242)
In(242)
) Ki, 2 1,
Remark 4.3. The O()) term hides an instance dependent factor. We will give its full specifies in the regret upper bound of Theorem 6.1.
K.a
Kia
+00,
K 0.
The bound in Theorem 4.2 shows that the total expected computation time of ATA remains within multiplicative factor of 1 + 4n In Bn of the optimal computation time Ck with additional remainder term that scales logarithmically with K. Since Cx = ?(K), ), this additive term is negligible compared to C* In practical scenarios, where computation time follows common distributions such as exponential or Gamma, the factor 1 is typically of order 1, and In B)) remains relatively small for the batch sizes commonly used in optimization algorithms like SGD.
The term conf(.,.) is derived from a known concentration inequality for sub-exponential variables with an Orlicz norm bounded by a (Lemma E in the Appendix).
Given the confidence bounds $ (s1, kn,k), the learner selects the action ag E A at round & that minimizes the loss É(-, sx), defined in (4). While nonconvex, we show in Appendix C that this optimization problem can be solved using a recursive routine, whose computational efficiency is On Inmin B, n)) min min(B,n) ))
Remark 4.1. Line 7 of the algorithm acts placeholder for the optimization method, where the optimization pa- rameters are updated using the quantities computed by the workers (c.g., gradients in the case of SGD). In this view, the allocation algorithm is independent of the specifics of
The reader might wonder if the more ambitious goal of de- riving bounds with a multiplicative factor of exactly is achievable. However, achieving this goal would require significantly more precise estimates of the expected compu- tation time [[C(a)) for all a E A. Since B(C(a)) depends
5

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
on the joint distribution of all workers in the support a, ob- taining such precise estimates would come at the cost of computational efficiency in the allocation strategy.
The expected total computation time CK of ATA-Empirical satisfies the same guarantee presented in Theorem 6.1, but we obtain an improved multiplicative factor of the additive logarithmic term. The precise expressions of these factors are provided in the next section, and they show that the guar- antees of ATA-Empirical adapt to the Orlicz norms || of each arm, while the guarantees of ATA depend on the maximum Orlicz norm max, ||
We note that it is unsurprising that 77 appears in the upper bound of Theorem 4.2, since having a heavier-tailed dis- tribution increases the gap between E(a, μ)) and E' [C(a]] through the convexity of C(-). Instead, the factor In In) arises because C()) is expressed as the maximum of up to B random variables. Moreover, in the edge cases where 0 (deterministic case) or B = (linear cost function), we guarantee that the expected computation time is at most an additive factor away from the optimal one.
6. Theoretical Results
In this section, we sketch the derivation of Theorem 4.2 for ATA and ATA-Empirical, through a regret analysis on the proxy losses. We define the expected cumulative regret of the proxy loss e(:; 11) after K rounds
5. Empirical Adaptive Task Allocation
The ATA procedure is based a lower confidence bound approach that relies on concentration inequalities. These bounds play a key role in performance, as sharper concen- tration bounds lead to more accurate estimates and reduce exploration of suboptimal options. Since workers' com- putation times follow sub-exponential distributions, their concentration behavior determined by the Orliez norm of the corresponding variables. In ATA, the only prior knowl- edge available is an upper bound the largest Orlicz norm among all arms. When the Orlicz norms of the arms' distri- butions vary significantly, this uniform bound may result in loose confidence intervals and inefficient exploration.
K
RK =
Σ EE [e(an, po) K KE(ā,, k=1
where ā E arg min ((a, 11) represents the optimal allo- cation over the workers. If multiple optimal actions exist, we consider the one returned by the optimization sub-routine used in ATA (line 5 of Algorithm 2).
We derive upper bounds the expected cumulative regret Rx- Based on these bounds, we provide the guarantees on the expected total computation time required to complete K iterations of the optimization process.
To address this issue, we introduce ATA-Empirical, which better adapts to the distribution of each arm. particularly its Orliez norm. This adaptation is achieved through novel data-dependent concentration inequality for sub-exponential variables. Unlike ATA, which depends on the maximum Or- liez norm, ATA-Empirical accounts for the individual Orlicz norms of all arms, denoted by (ax) This improvement reflected in the upper bounds on regret presented in Sec- tion 6. In practice, this leads improved performance at least some settings, as shown in our simulations in Section 7. However, this increased adaptivity comes with trade-off since ATA Empiricll requires an upper bound on the quantity max; a 14, rather than a bound on the largest Orlicz norm. That said, for many distributions of interest, the ratios az/Ms across different arms tend to be of the same order, whereas their Orlicz norms can vary significantly.
6.1. Guarantees for ATA
For each worker i [m], recall that a, denote the prescribed allocation of the optimal action ā. Define k; as the smallest integer satisfying
(a. + k:) > l(a,
(10)
From the definition above, it follows that if the learner plays an action ag at round h such that ag.; 2 ā; + k, then E(ax, 11) > E(a, 14). Thus, k; can be interpreted as the smallest number of additional units allocated to worker i that result in a suboptimal loss. Moreover, for every worker i e [n], we have kg E {1,2] (see Lemma D. in the Appendix).
The next result provides upper bound on the expected regret of ATA.
Theorem 6.1 (Proof in Appendix D.I). Suppose that As- sumption 3.1 holds. Then, the expected regret of ATA with inputs (x)x) satisfies
The ATA-Empirical procedure differs from ATA only in the lower confidence bounds it uses. These bounds are derived from the novel concentration inequality in Lemma 6.2 and are defined for arm E [n]] at round & E [K] as
RK S 2n max Bu, [(a. sE(n]
$ Pi,k 1 1- 2n (v
ln(2k2)
In 12k2
a² k:) E(ā, 11) In K,
(8)
((o: + ki)j4 l(ā H1)2
Kip
Ki,
where a 2 max
||x - Holler and c is a numerical
where 7) = max;e(n) ai/M3-
constant.
6

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
The first term in the regret upper bound is independent on the number of rounds K. The second term, however, grows logarithmically with K, which aligns with the behavior observed in stochastic bandit problems in the literature.
accumulate enough samples to stop over-allocating worker can be upper-bounded by
of + k;) ((ā: + ki)su 1))²
In the case where B = 1, our setting reduces to the problem of regret minimization for the standard multi-armed bandits. Observe that in this case E(ā. μi) minke[n] (4g, & for all i E [n]. Therefore. the guarantees of Theorem 6.1 recover the known optimal bound
In the regret bound of Theorem 6.1, the term a² appears in- stead of af because the learner's prior knowledge limited to an upper bound a > max; || - 14ll on the maxi- mal Orliez norm of the arm distributions. Finally, consid- ering that the worst-case excess loss incurred when over- allocating worker is Bp; E(ā, μu), we obtain the stated bound.
o Σ [i((K)/41
of the standard MAB setting, where A := p4 ming Uz-
6.2. Guarantees for ATA- Emppricaa
Proof sketch. In standard and combinatorial MAB prob- lems, regret bounds are typically derived by controlling the number of rounds in which the learner selects subopti- mal arms. These bounds are often of the order In (n(K)//² A where A denotes the suboptimality gap and quantifies the exploration cost required to distinguish optimal actions from suboptimal ones.
We present theoretical guarantees for ATA-Empirical by pro- viding an upper bound on the expected cumulative regret (9), As discussed in Section 4, ATA-Empirical leverages lower confidence bounds derived from a novel data-dependent concentration inequality introduced below, The proof of this result is detailed in Appendix E.
: Xn be i.i.d. positive random vari- ables with mean ps, such that a [ plle, < +o0. Let Xn denote the empirical mean. For 8 E (0, 1), let
Lemma 6.2. Let X1....;
In our setting, the problem is more complex since the learner must only choose which arms to pull but also determine the allocation of resources across selected arms. With this in mind, we develop the following key arguments leading to the bound in Theorem 6.1.
log log(2/0) 1og(2/6)
Cn,6: = 2/
n
We define over-allocation for worker at round as the event where 0; 2 a; + kū By definition of k; (see (10)), this implies that f(ax, (1) > f(a. We define a bad round round where E(az, (1) > E(a, (b), and we say that a bad round is triggered by arm i when di,a = E(as, 14) > ((ā 1t). Then, the proof revolves around establishing an upper bound on the total number of bad rounds.
where 17 = aa/, Then, with probability at least 1 8, we have
# X, (1 n/On.5)+
Moreover, if nn, 4. then, have with probability at least 1 6, we have
To derive this bound, we consider the number of samples required to verify that the mean computation time of worker i under over-allocation exceeds the optimal waiting time f(a ((, p1). Specifically, we need to test whether the mean of the corresponding distribution, at least (ā, +i,),, surpasses the threshold f(ā. pt). This is equivalent to testing whether
(1 4 CnC (1- spsss Su (1+ [1+3)Cnnd
Using the concentration inequality above, construct the lower confidence bounds & defined in (8). The fol- lowing theorem provides an upper bound on the regret of Emppriaiii
E(a. 1L) a4
p4
Theorem 6.3 (Proof in Appendix D.2). Suppose that As- sumption 3.1 holds. Then, the expected regret of ATA- Empirical with inputs (B, 1), satisfies
Using the concentration inequality applied in our analysis, the number of samples required for this test is of the order:
R S 2n max BBp A(ā, 11))
af (1:
)
iE(n]
(11)
[(a, +,, 1))=
( + kz) Ela, u))
In K, ((ā; + kz) (1)²
During rounds where worker is over-allocated, the learner collects at least a + k2 samples from the corresponding dis- tribution. Therefore, the total number of rounds required to
1=1
where 1) 2 maxie(n) 0s/14 and e numerical constant.
7

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
Comparing the bounds for ATA-Empirical and ATA, ob- serve key differences. Unlike the bound in Theorem 6. which incurs squared maximal Orlicz norm penalty of a² for all terms in the upper bound, ATA-Empirical benefits from its adaptive nature, leading to a term-specific factor of s7² 1u.- In the case where the arm distributions have a ratio ax ju; of the same order (such as the exponential distributions), the bound of Theorem 6.3 shows that ATA-Empirical adapts to the quantities ag as we have, in the last case, TT14; 0zz-
in (9) in each iteration, and UTA (Uniform Task Allocation), which distributes B tasks uniformly among the n workers. If n > B. then in UTA we select B workers at random, each one tasked to calculate one stochastic gradient. Our algorithms aim to achieve a performance close to the one of OFTA, without any prior knowledge of the true means.
For ATA we set a an 4-29/ni, while for ATA-Empirical we use 1j = 1. The results of our experiments are shown in Figure 1. As expected, GTA is the fastest in terms of runtime (first column), but it performs poorly in terms of total worker time (second column). This is because it uses all devices, most of which perform useless computations that are never used, leading to worse performance as the number of workers increases. In fact, its performance can become arbitrarily worse. On the other hand, OFTA performs best in terms of total worker time. Although it is slower in terms of runtime, the difference is by a constant factor that does not increase as grows. This is because additional workers are less efficient and do not provide significant benefits for GTA.
7. Experiments
In this section, we validate our algorithms by simulating a scenario with n workers, where we solve simple problem using SGD. In each iteration, we collect B = 23 gradients from the workers and perform gradient descent step.
The objective function f Rd + R is convex quadratic defined as
f(z) = 2 x Ax - b bTx,
where
Turning our attention to our algorithms, both ATA and ATA- Empirical initially behave like UTA, it is expected by the need to perform an initial exploration phase with uniform al- locations. However, after this phase, they begin to converge to the performance of OFTA.
2
-1
0
A
1
E RR dxd
4
-1
0
2
The last two columns contain plots that confirm our theo- retical derivations. The third plot validates Theorem 4.2, showing that ATA and ATA-Empirical converge to OFTA up to constant. The final column shows the averaged cumulative regret, vanishing over time as predicted by Theorems 6. and 6.3.
--1
1
0
b
E IRd
4
We denote f as the minimum value of the function f. Each of the n workers is able to calculate unbiased stochastic gradients g (ar) that satisfy
Table 1: Ratios of total worker times and runtimes required to achieve f(x) - f < 10-5 For total worker time, we divide the total worker time of GTA by the corresponding total worker times of the other algorithms listed. For run- time, we do the opposite, dividing the runtime of the other algorithms by the runtime of GTA, since GTA is the fastest. To simplify the naming, we refer to ATA-Empirical as ATA-E.
E
- - < 0.012.
This is achieved by adding Gaussian noise to the gradients of f.
The computation time for worker is modeled by the distri- bution
TOT. WORKER TIME RATIO
RUNTIME RATIO
n
= 29 + Exp (29)i).
ATA
ATA-E
OFTA
ATA
ATA-E
OFTA
17
1.3
1.26
1.26
1.73
1.75
1.74
for all i E [n], where Exp(3) denotes the exponential dis- tribution with scale parameter 3. The expected value this distribution is μ4 = 2- 29. Furthermore, the Orliez norm satisfies the bound a, < 2p4-
51
2.91
2.69
3.03
2.43
2.45
2.17
153
7.22
7.02
9.1
3.41
3.14
2.17
459
12.45
14.1
27.3
6.36
5.51
2.17
We consider three benchmark algorithms. GTA-SGD, orig- inally introduced as Rennala SGD by Tyurin & Richtárik (2024). Additionally, we include OFTA (Optimal Fixed Task Allocation), which assumes the oracle knowledge of the mean computation times and uses the optimal allocation à
In Table 1, we compare the results numerically. Both the total worker time ratio and runtime ratio increase as n grows. The total worker time ratio increases because GTA becomes less efficient, using more resources than necessary. The nuntime ratio grows for ATA and ATA- ETA-iiiiiical since a larger
8

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
50"
507
13300
101
03
200
0.0
13
107
1200
30"
330
3.
100
Rantims
0.
Tea nottar
2.0
1/15
0.3
1.0
0.0
0.5
35
10
50-
1400
430
4.50
30*
234
1408 wobkr
1812
13
101
Tantrm
T5ta Nother
2.122
Figure 1: Each row increases the number workers by a factor of 3, starting from 17, that is, n = 17,51, 153, 459 from top to bottom. The first column shows runtime vs. suboptimality. The second column also plots suboptimality, but against total worker time, i.e., >i a in Algorithm 1. The third column presents the average iteration time, given by Cz/k over all iterations k. The last column displays the averaged cumulative regret, as defined in (9).
number of workers requires more exploration. However, for OFTA this ratio remains unchanged, as discussed earlier.
on the CIFAR-100 dataset (Krizhevsky et al., 2009) using Adam (Kingma & Ba, 2014).
We remark that in these experiments we started all runs for ATA and ATA Empiricaiaal without prior knowledge of the computation time distribution. However, in real systems, where these algorithms used multiple times, prior esti- mates of computation times from previous runs could be available. With this information, ATA and ATA-Empirical would be much faster, as they would spend less time on ex- ploration, approaching the performance of OFTA in a faster way. We validate this through experiments presented in Appendix A.S.
A 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) Center of Excellence for Generative AI, under award number 5940, iii) SDAIA-KAUST Center of Excel- Ience in Artificial Intelligence and Data Science.
Impact Statement
In Appendix A.I, we conducted similar experiments with different time distribution, where the mean times vary linearly across the arms In Appendix A.2, we examine sce- marios with varying client time distributions. Additionally, in Appendix A.3, we analyze regret performance, confirm- ing its logarithmic behavior as predicted by Theorems 6.1 and 6.3, Finally, in Appendix A.4, we trained a simple CNN
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
9

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
References
Kairouz, P., McMahan, H. B., Avent, B., Bellet, A., Bennis, M., Bhagoji, A. N., Bonawitz, K., Charles, Z.,, Cormode, G., Cummings, R., et al. Advances and open problems in federated learning. Foundations and Trends in Machine Learning, 14(1 14(1-2):1-210, 2021.
Auer, P. Finite-time analysis of the multiarmed bandit prob- lem, 2002.
Boucheron, S., Bousquet, O., Lugosi, G., and Massart, P. Moment inequalities for functions of independent random variables. 2005.
Kingma, D. P. and Ba, Adam: A method for stochastic optimization. arXiv preprint 14Xix::412.69880 2014.
Cesa-Bianchi, N. and Lugosi, G. Combinatorial bandits. Journal of Computer and System Sciences, 1404 1422, 2012.
Konečný, J., McMahan, H. B., Yu, X., Richtarik, P.,
Suresh, A. T., and Bacon, D. Federated learning: Strate- gies for improving communication efficiency. In NIPS Workshop on Private Multi-Party Machine Learning. 2016.
Chen, J., Pan, X., Monga, R., Bengio, S., and Jozefowicz, R. Revisiting distributed synchronous SGD, arXiv preprint arXiv: 1604. 00981, 2016a.
Krizhevsky, A., Hinton, G., et al. Learning multiple layers features from tiny images. Technical report, University of Toronto, Toronto, 2009.
Chen, W. Wang, Y., and Yuan, Y. Combinatorial multi- armed bandit: General framework and applications. In International Conference on Machine Learning, pp. 151- 159. PMLR, 2013.
Kveton, B., Wen, Z., Ashkan, A., and Szepesvari, C. Tight regret bounds for stochastic combinatorial semi-bandits. In Artificial Intelligence and Statistics, pp. 535-543. PMLR, 2015.
Chen, W., Hu, W., Li, F., J., Liu, Y., and Lu, P. Combina- torial multi-armed bandit with general reward functions. Advances in Neural Information Processing Systems, 29, 2016b.
Lattimore, T. and Szepesvári, C. Bandit algorithms. Cam- bridge University Press, 2020.
Combes, R., Talebi, M. S., Proutiere, A., and Lelarge, M. Combinatorial bandits revisited. Advances in Neural Information Processing Systems, 28, 2015.
Li, T., Sahu, A. K., Zaheer, M., Sanjabi, M., Talwalkar, A., and Smith, Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Systems, 2:429 450, 2020.
Dean, J. and Barroso, L. A. The tail at scale. Communica- fions of the ACM, 56(2):74-80, 2013.
Lin, T., Li, J., and Chen, W. Stochastic online greedy learning with semi-bandit feedbacks. Advances in Neural Information Processing Systems, 28, 2015.
Dutta, S., Joshi, G., Ghosh, S., Dube, P., and Nagpurkar, P. Slow and stale gradients win the race: Error- runtime trade-offs in distributed SGD. In International Conference on Artificial Intelligence and Statistics, pp. 803-812. PMLR, 2018
Maranjyan, A., Safaryan, M. and Richtárik, P. Grad- Skip: Communication accelerated local gradient meth- ods with better computational complexity. arXiv preprint arXiv:2210.164400 2022.
Gai, Y., Krishnamachari, B., and Jain, R. Combinatorial network optimization with unknown variables: Multi- armed bandits with linear rewards and individual obser- vations. IEEE/ACM Transactions on Networking, 20(5): 1466-1478, 2012.
Maranjyan, A., Omar, O. S., and Richtárik, P. MindFlayer SGD: Efficient parallel SGD in the presence of heteroge- neous and random worker compute times. In Conference on Uncertainty in Artificial Intelligence, 2025a.
Maranjyan, A., Tyurin, A., and Richtárik, Ringmaster ASGD: The first Asynchronous SGD with optimal time complexity. In International Conference on Machine Learning. 2025b.
Gelenbe, E. and Mitrani, I. Analysis and synthesis of com- puter systems, volume 4. World Scientific, 2010.
Gross, D., Shortle, J. F., Thompson, J. M., and Harris, M. Fundamentals of queueing theory, volume 627. John wiley& & sons, 2011.
Maurer, A. and Pontil, M. Concentration inequalities un- der sub- Sab-Gaussian and sub-exponential conditions. Ad- vances in Neural Information Processing Systems, 34: 7588-7597, 2021.
Hadjis, S., Zhang, C., Mitliagkas, I., Iter, D., and Ré, C. Omnivore: An optimizer for dulti-device deep learning on CPUs and GPUs. arXiv preprint arXiv: 160.066400 2016.
McMahan, B., Moore, E., Ramage, D., Hampson, S., and y Arcas, B. Communication-efficieee learning deep
10

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
networks from decentralized data. In Artificial Intelli- gence and Statistics, pp. 1273-1282. PMLR, 2017.
McMahan, H. B., Moore, E., Ramage, D., and y Arcas, B. A. Federated learning of deep networks using model averaging. arXiv preprint 1rXiv:1602.05622 2:2, 2016.
Mitliagkas, I., Zhang, C., Hadjis, S., and Ré, C. Asynchrony begets momentum, with an application to deep learning. In 2016 54th Annual Allerton Conference on Communi- cation, Control, and Computing (Allerton), pp. 997-1004. IEEE, 2016.
Nguyen, J., Malik, K., Zhan, H., Yousefpour, A., Rabbat, M., Malek, M., and Huba, D. Federated learning with buffered asynchronous aggregation. In International Con- ference on Artificial Intelligence and Statistics, pp. 3581- 3607. PMLR. 2022.
Tyurin, A. and Richtárik, P. Optimal time complexities of parallel stochastic optimization methods under fixed computation model. Advances in Neural Information Processing Systems, 36, 2024.
Wang, S. and Chen, W. Thompson sampling for combi- natorial semi-bandits. In International Conference on Machine Learning, pp. 5114-5122, PMLR, 2018.

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
5400
801
2200
0.0
137
3.50
30m
Rantma
Tit nottardr
1/42
0.5
1.0
0.0
0.5
35
33
10
11300
7000
301
1012
180
Figure 2: We use the same setup as Figure I, with each row tripling the number of workers, starting from n = 17.
A. Additional Experiments
The objective function is a convex quadratic function f Rd + IR defined as
f(x: 2 A.r bTa: x,
where
2
-1
0
A
-1
RR ddd
1 and b
0
E Rdd
4
-1
We denote f as the minimum value of the function f. Each of the workers is able to calculate unbiased stochastic gradients (zx) that satisfy
[[(()) (f())l s 0.01
This is achieved by adding Gaussian noise to the gradients of f.
The experiments were implemented in Python. The distributed environment was emulated machines with Intel(R)
Xeon(R) Gold 6248 CPU @ 2.50GHz.
A.1. Linear Noise
In this section model the computation time for worker i by the distribution
2; = 29i + Exp() for all i E [m]
The expected value this distribution 14 = 2 29i. Furthermore, the Orlicz norn satisfies the bound a; S 2/4
We again set B 23 and run simulations similar those Section 7. The results are shown Figure 2.
12

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
The important difference the previous Figure is that here ATA-Empirical performs better than ATA. This is because the Orlicz norm 4 29n is much larger.
Similarly, we provide numerical comparison in Table 2.
Table 2: This table presents ratios similar to those in Table 1.
TOTAL WORKER TIME RATIO
RUNTIME RATIO
π
ATA
ATA-Empirical
OFTA
ATA
ATA-Empirical
OFTA
17
2.32
1.91
2.1
1.71
1.71
1.58
6.71
5.02
6.29
3.27
2.12
1.58
153
3.41
8.68
18.87
7.96
4.5
1.58
A.2. Heterogeneous Time Distributions
So far, we have only considered cases where clients follow the same distributions but with different means. In this section, we extend our experiments to cases where the distributions themselves differ. We consider five distributions: Exponential, Uniform, Half-Normal, Lognormal, and Gamma. We group five workers with these five distributions that each group has the same mean, then vary the mean across different groups. More concretely, use:
Exp(c(5g + 1)),
Uniform
115 2
I (0, c(5g + 1) 1.
Lognormal (log(c(5g + 1))/2, Vlog(c(5g+ + 1)))
Gammaa (c(5g+ 1))², (56 with shape and scale parameters.
Next, we add a constant c(5g + 1) to all the distributions, where = 29, and represents the group number, starting from 0. The clients are divided into n/5 groups
The results of the experiments are shown in Figure The plots demonstrate that the algorithms are robust across different distributions.
A.3. Regret
In this section, we verify Theorems 6.1 and 6.3 on regret through simulations. We set 20 and computation time for worker i following the distribution
5, with the
V; = Exp(2i), for all i E [m]]
We ran the simulation five times, and the plots include standard deviation bars, although they are not visible. The results presented in Figure 4.
As expected, the regret grows logarithmically.
A.4. Real Dataset
In this section, we present an experiment where we train convolutional neural network (CNN) the CIFAR-100 dataset (Krizhevsky et al., 2009). The network consists of three convolutional layers and two fully connected layers, with total of 160k parameters.
We use the Adam optimizer (Kingma & Ba, 2014) with a constant step size of 8 1-105 The computation time of the workers follows the same setup as in Figure 2. The results are shown in Figure
13

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
1000
301
307
0.0
13
107
3250
1750
7030
30m
Santma
Tta natarrdaa
1-42
65
1.0
0.0
0.5
35
10
2030
10300
13300
30"
1012
180
Figure 3: Each row corresponds to an increasing number of workers, with n= 15, 45, 150 from top to bottom. We consider five distributions--Expoeee Exponential, Uniform, Half-Normal, Lognormal, and Gamma Ganma-ggouping them have the same mean and then varying the mean across different groups. The results demonstrate that the algorithms remain robust across different distributions. The columns represent the same as in Figure 1.
A.5. Impact of Prior Knowledge on Time Distributions
In real-world systems where multiple machine learning models are trained, estimates of computation times from previous runs may be available. With this prior knowledge, ATA and ATA-Empirical can be much faster, as they spend less time on exploration and quickly approach the performance of OFTA.
To illustrate this, we vary the number of prior runs, P. Since our algorithms operate independently of the underlying optimization process, we first focus solely on the bandit component, updating the confidence scores of machines over several iterations. We then apply the loss curves different segments of the bandit phase and compare the results P increases. larger P yields more accurate estimates.
The optimization setup remains the same in Figure 5, with B 23 and n 51. The results presented in Figure 6.
B. Concrete Optimization Methods
In this section, we provide concrete examples of optimization algorithms using the ATA and GTA allocation strategies.
For optimization problems, we focus on SGD and Asynchronous SGD. Other methods, such as stochastic proximal point methods and higher-order methods, can be developed in a similar fashion.
B.1. Stochastic Gradient Descent
For SGD, is important distinguish homogeneous and heterogencous cases. Let start from the homogeneous case.
B.1.1. HOMOGENEOUS REGIME
Consider the problem of finding an approximate stationary point of the optimization problem
min ZER
{f(x)
EE [f(x;))]]
(12)
14

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
ATA: Empincal
Regret
0.0
0.5
1.0
1.5
2.0 Iterations
2.5
3.0
3.5
4.0
1e6
Figure 4: Regret growth over iterations.
We assume that each worker is able to compute stochastic gradient f(x;; {) satisfying RELD [1] {) [(())]]]] S o² for all x E Rd.
In this case, SGD with allocation budget B becomes Minibatch SGD with batch size B. The next step is determining how the batch is collected. For ATA, we refer to this method as SGD-ATA, as described in Algorithm 2.
Algorithm 2 SGD-ATA (Homogeneous)
1: Optimization inputs: initial point æ% Rd, stepsize 7 >> 0
2: Allocation inputs:
allocation budget B
3: Initialize: empirical means f4,1 = 0, usage counts K; = 0, and usage times Tù O, for all i E [n]
4: for k = 1,...... K do
5: Compute LCBs (s),,) for all E [n]
6: Find allocation: ag E arg min E(a, (x))..). Allocate a₁ tasks to each worker E [n]
8: Update
E
0k+1
B
jml
9:
for i such that 0i,k * 0 do
10:
K2 = KiA + 0i,k
T Ti, 2..6 x()
(1,t+1 K,
1
13: end for
14: end for
In this case, each task consists in calculating the gradient using the device's 's local data, which is assumed to have the same distribution as the data on all other devices. Because of this, it does not matter which device performs the task. The method then averages these gradients to obtain an unbiased gradient estimator and performs a gradient descent step.
Now, let give the version of Minibatch SGD using greedy allocation Algorithm 3.
15

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
4.0
XX
4.0
3.5
3.5
optnm
2.0
3.0
2.5
0.0
0.0
0.2
0.4
0.6
0.8
1.0
1.2
1.6
Runkima
4.0
Optrmal
1735
2.5
2.00
22.0
1.5
1.5
1.0
1.0
0.5
1.5
2.0
2.5
1.0 Total worker time
Runtime
160
1610
22
3.5
0th
3.5
2.0
3.0
2.5
2.5
2.0
220
1.5
1c10
Figure 5: We use the CIFAR-100 dataset (Krizhevsky al., 2009). The model a CNN with three convolutional layers and two fully connected layers, totaling 160k parameters. The Adam optimizer (Kingma & Ba, 2014) is used with constant step size of 8 10 The computation time of the workers follows the same setup in Figure 2, where the mean time increases linearly. The batch size remains the same at B = 23. Each row corresponds to different number of workers, with n 17, 51, 153 from top to bottom.
Algorithm SGD-GTA (Homogeneous)
1: Input: initial point æn E IRd, stepsize > 0, allocation budget B
2: for k = 1...,,, K do
3: b=0
4: Query single gradient from each worker i [m]]
5: while b < do
6: Gradient V far Ev.) arrives from worker iks
7:
9k = 9k + f Ex b = b + 1
8: Query gradient at from worker iR,
9: end while
10: Update the point: Tk+1 36 7 g"
11: end for
Algorithm is exactly Rennala SGD method proposed by Tyurin & Richtárik (2024), which has optimal time complexity when the objective function is non-convex and smooth.
If the computation times are deterministic, then GTA makes the same allocation in each iteration. In that case, SGD-ATA will
16

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
4.0
XX
4.0
3.5
3.5
2.0
3.0
2.5
0.0
0.5
1.0
1.5
2.0
2.5
3.0
3.5
0.0
0.5
1.0
1.5
2.0
Optrmal
1725
22.0
1.5
1.5
1.0
1.0
0.5
1.5
2.0
Runtime
160
Total worker time
1610
3.5
695
2.0
3.0
2.5
2.5
2.0
220
1.5
2.0
2.5
Figure 6: We use the same optimization setup as in Figure with B 23 and 51. The number of prior iterations. P. varies across rows, starting from the top with P 0,5 10 5 106. As the number of prior iterations increases, observe that the training of ATA and ATA-Empirical accelerates, bringing their performance closer to the optimal performance of OFTA.
converge this fixed allocation. If the times random, the allocation found by GTA may vary in each iteration. In this case, SGD-ATA will approach the best allocation for the expected times.
B. 1.2. HETEROGENEOUS REGIME
Now let us consider the following heterogeneous problem
[(x) = 1 D Ee, ~D Ui(a:
min #ERd
Here each worker has its own data distribution D,. We start with the greedy allocation. The algorithm is presented Algorithm 4.
17

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
Algorithm 4 SGD-GTA (Heterogeneous)
1: Input: initial point zo E IR",, stepsize 7 > 0, parameter S
2: for = 1..... K do
3: S; 0 and gi.k = 0 for all i [n]
4: Query single gradient from each worker i [n]
5: while -.
do
6: Gradient (aps Ek) arrives from worker j
9j.k = 9j.k + V ax: Ex): 8j = 8j +
8: Query gradient at æ from worker )
9: end while
10: Update the point: 2k+1 7/2 Σi Ki.k
11: end for
Algorithm 5 presents the Malenia SGD algorithm, proposed by Tyurin & Richtárik (2024), which is also optimal for non-convex smooth functions.
In each iteration, Algorithm 4 receives least one gradient from each worker. Building this idea, design method incorporating ATA, given in Algorithm 5.
Algorithm 5 SGD-ATA (Heterogeneous)
Optimization inputs: initial point æ Rd, stepsize 7 >> 0
2: Allocation inputs: allocation budget B
3: Initialize: empirical means fr,1 = 0, usage counts K2, = 0, and usage times Tù, O, for all i e [n]
4: for k = l..... K do
5: Compute LCBs (s),,) for all E [n]
6: Find allocation: ag = RAS( B)
7: Allocate 0, + tasks each worker [n]
8: Update æ::
7 72
G1a+1
Øk+1
V
71
0i,k +
i=1
j=1
9: For all é E [n]. update:
Kip+1 = Ki,h + ai.k
0i
Tekt1 Ti, +
i.k
1=1
Tik+1
Pa,41 Ki,41
end for
B.2. Asynchronous SGD
Here, focus on the homogeneous problem given in Equation (12). The greedy variant, Ringmaster ASGD, proposed by Maranjyan et al. (2025b) and, like Rennala SGD, achieves the best runtime.
We now present its version with ATA, given in Algorithm 6.
Here, the task remains gradient computation, but each worker's subsequent tasks use different points for computing the gradient. These points depend on the actual computation times and the asynchronous nature of the method, hence the name
18

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
Algorithm 6 ASGDD-TA
Algorithm 7 ASGD
1: Optimization inputs: initial point æo E R stepsize
I: Input: Initial point 2o E R stepsize y > 0, allocation vector with ||a]]] B
y > 0
2: Allocation inputs:
allocation budget B
2: Workers with d; 0 start computing stochastic gradi- ents at 20
3: Initialize: empirical means fi,1 0, usage counts
K = 0, and usage times T;, = 0, for all E [n] 4: for k = -....... K do
3: for s = 0....., B - 1 do
4: Receive gradient V f S8+6. from worker
5: Compute LCBs (s).,) for all E [n]
5: Update: Fz+1 X. >V ( E848,)
6: Find allocation: ag = (*s B)
6: if a; then
7: Update z& using Algorithm with allocation as
7: Worker i begins computing Vi(
8: For all i such that ai # 0, update:
8: Decrease remaining allocation for worker by one:
a: =
KiA+1 Ki, + di,k
9: end if
aikk
10: end for
x
Tix+1
Ti,k +
11: return: BB
Ta,+1
The sequence {6.) represents delays, where 6, 2 is the difference between the iteration when worker started computing the gradient and iteration s, when was applied.
Ai,+1 Kg,+1
9: end for
Asynchronous SGD.
C. Recursive Allocation Selection Algorithm
In this section, we introduce an efficient method for finding the best allocation. Given LCBs sy and allocation budget B, each iteration of ATA (Algorithm 1) determines the allocation by solving
aj E arg min l(a, 84),
aEA
where
E(a, (4) = max 0.44 = || @ polooo? (E
with @ denoting the element-wise product. When clear from context, we write f(a) instead of E(a, 14).
In the early iterations, when some 8; values are 0, ATA allocates uniformly across these arms until all 8; values become positive. After that, the allocation determined using the recursive routine in Algorithm Remark C.I. The iteration complexity of RAS (n in(min) ( B, n)) + min B, n)² In fact, the first n In(min (B, n)) term arises from identifying the smallest B scores. For the second term, note that in (13), we have r S nn..
C.1. Optimality
We now prove that RAS finds the optimal allocation, as stated in the following lemma. Lemma C.2. For positive scores
$1 S 82 S ...... .. Sn, RAS (Algorithm 8) finds an optimal allocation h A, satisfying h arg min ||a sllaa aeA
Proof. We prove the claim by induction the allocation budget B.
Base Case (B = 1): When B = 1, RAS (Algorithm 8) allocates the task to worker with the smallest score (line 9). Thus, the base case holds.
Inductive Step: Assume that RAS finds an optimal allocation for budget 1, denoted by
h RAS( (81, .1...., Sn; B 1).
19

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
Algorithm 8 Recursive Allocation Selection (RAS)
1: Input: Scores .1.....8m Sn, allocation budget B
2: Assume without loss of generality that 81 82 S ... S (i.e., sort the scores)
3: if' B = then
4: return: (1,0.....0)
5: end if
6: Find the previous best allocation:
a = (d.,...;.p)) = RAS (s), (s1,..., 8nj B 1)
7: Determine the first zero allocation:
minfi a, 0}, if a₂ =0 otherwise
r
(13)
8: Find the best next query allocation set:
M argmin KE|r]
[I(a + e;) o $|1
where e; is the unit vector in direction i. 9: Select j E M such that the cardinality
arg max (a; ej,)ss iE[r]
is minimized return: a + Ej
We need prove that the solution returned for budget B, denoted by h h er, is also optimal.
Assume, for contradiction, that there exists E A such that + h and E(a) E(h). Write ā + Ey for some & [m]. Observe that ||a|| B - because E A.
We consider two cases based on the value of E (h + ez)::
l [( + e,) hg for some & + r. this case, adding one unit index r does not change the maximum value, i.e.,
e (h) = f {h + e,). By the inductive hypothesis, minimizes {(x) for budget B 1. Therefore, we have
((a) 2 l (a) 2 (h) (h + +e,) = E(h),
which contradicts the assumption that (aa) < E(h).
e (h + e,) (h+ 1) 8r. By the algorithm's ' logic, (h, + 1) 8, S (h; + 1)s; s; for all i i Since {(h + e,) S f(h + ey) and we assumed l(a + ee)) E(a) < ((h) e(h + e,), then à & h otherwise E(ā + eq) < f(ā + ey).
Given that ||a|] = | this implies that there exists some u E [n]] such that 0 < ā. S Th and another index V E [n]] where a, 2 h2 + 1.
In addition, note that r is chosen such that (h + e,) is minimum. Using the fact that & (h + e,) = (h, + 1) Sr, we have that for any index we also necessarily have (h + eg) = (hg + 1) Sq. Using this, we deduce
((h) (h ex) s se(hte)) (Fr 1) st max ā,s, l (a) - l(a),
where in the second inequality we used the fact that a 2 h 1 and in the last incquality we used the fact that the loss is not decreasing for add one element to the vector. This chain of inequalities again contradicts the assumption that f(a) {h)).
Since both cases lead to contradictions, we conclude that no a E exists with E(a) < E(h). Thus, RAS produces an optimal allocation for budget B.
20

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
C.2. Minimal Cardinality
Among all possible allocations RAS choose one that always minimizes the cardinality of the set:
arg max 0,Si sE(n)
The reason for this choice is just technical as it allows the Lemma D.I to be true. Lemma C.3. The output of RAS ensures the smallest cardinality of the set:
arg max 0;5; *E(n]
among all the optimal allocations a.
Proof. This proof uses similar reasoning as the one before.
Let h (s s; B), and denote the cardinality of the set arg maX 0,8; for allocation a by
Cs(a)= = arg max 0:84 21.
xc(m]
We prove the claim by induction on B.
Base Case (B = 1): For = 1, there is a single coordinate allocation, thus Ci (h) = 1, which is the smallest possible cardinality.
Inductive Step: Assume that RAS finds an optimal allocation for budget with the smallest cardinality, denote its output by
h RAS(S1) ......., Sn; B 1).
We need to prove that h + e, minimizes CB(a) among all optimal allocations for budget B.
Assume, for contradiction, that there exists such that a + h, E(a) E(h). and Cs(a) < Cb(h). Write a t cg for some q E [n]. We consider three cases:
Cp (h) 1. Since the minimum cardinality is exactly must have Cn(()) 2 Cn(h), that contradicts assumption.
Cn(h) = Cn-1 h) > 1. This occurs when ((h) f (h) + (hr + 1) 8r. By the optimality of h, we have (h) (a) < {(a) {(h) E(h). which implies (ā) {(a). Therefore, Ca-1 (α) S Cn(a). Since the induction hypothesis holds for B we have CB-1 (h) S CB-1 (a). Thus,
(h) CB-1 (h) < CB-1 1 (a) s Cp(a),
which leads to a contradiction.
(nhh) Cn-1 (h) This occurs when {(h) (h) (h,, 1) 8y. Proceeding in the previous case, have (()) f(a), and hence CB-1 (a) < Ca(a). Since the induction hypothesis holds for B - 1. know
CB-1 (h) < CB-1 (@).
We now have additional cases:
If CB CB-1 (h) + 11 then
Cz(h) CB-1 (h) + +1= CB-1 Cz(a),
which leads a contradiction.
21

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
Now assume CB-1 (a) CB-1 (h). We will show that in this case, CB(a) CB-1 (a) By contradiction, suppose (a(a) Cb-1 (a), which implies (āg 1)8g < {(a)) Let & be index such that 0g8k {(a).
Construct new allocation a' a eg Ek. Then,
Cn-1 (a') CB-1 (a) - CB-1 (h).
which contradicts the induction hypothesis. Thus, C3(a) CB-1 (a) Using this, we have
CB(h)= = CB-1 (f) CB-1 (a) + = CB(a),
which again contradicts Ca(a) < (b(h)
This concludes the proof.
D. Proofs of Theorem 6.1, Theorem 6.3, and Theorem 4.2
We start by recalling the notation. For E [n] and E [K], (x(*) distribution V4. When using an allocation vector a₃ E A, the total computation time of worker at round & Ea x(x) when dg.k > l = ( jK) the vector of means. For each [K], when using the allocation vector ag, we recall the definition of the proxy loss f: A 12 120 by
1,8 )uEBB) denote independent samples at round from
(a A) = max di
(E(m]
where À = (Ai,..., , An) input of For each A, let ax E A, be the action minimizing this loss
is a vector of non-negative components. When À = μe, we drop the dependence on the second
a E mrg f(a, A)
aEA
We drop the dependency on from ap to ease notation. The actual (random) computation time at round is denoted by
C' A
asSk
Class) == max H
(u)
(14)
1
Let a" be the action minimizing the expected time
a' E arg min E [[((a)]
aeA
The expected regret after K rounds is defined as follows
K
RK = Σ -((aa))
For the remainder of this analysis we consider a E arg min {(a) found using the RAS procedure. For each E [n]. recall that k the smallest integer such that
(a, + k4j4 {(a).
(15)
Below present technical lemma used in the proofs of Theorems 6.1 and 6.3.
Lemma D.I. Let i (I1, .i,,..,,.. ...,,,)) E Ro Let a the output of RAS(ar; B). For each i, E [m], we have
a,m, S (a; + 24.
Proof. Fix x E R2. and let a = RAS(x;; B). The result is straightforward when min x; = 0. vE(n)
22

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
Suppose that x; 0 for all [n]. Let 2 denote the cardinality
arg max 0,20 iE[n]
Fix i.j E [n], let & E arg maxie[1] Og3i, We need to show that
ag S (a; + 1)
We a proof by contradiction. Suppose that have agTk (a: + 1),, consider the allocation vector a' A given by de
ag 1, d a, + 1 and d' a. when k) Let R == max Z.. We have
l(a', .z) max dh max ((as + 1) (ax Jar.. R)
wE(n]
We consider two cases:
Suppose that & = (i.e., the only element in [nx such that og = l(a, a)) is k), then we have necessarily R < agTb Morcover, by the contradiction hypothesis, (a; + 1)x; < agTk. Therefore.
fla', x) max{(a; + 1)24, (ax 1)rk, R} agzz E(a,
which contradicts the definition of a.
Suppose that 8 2 2, since by hypothesis ag.r > (a; 1)x,, we clearly have a,æ; < E(a, x) therefore among the set [n]] \ (k, i) there are exactly elements such that ag t(a, x) In particular, this gives
lla' x)) max {(a; + 1) (az 1) Dr R R E(a. x) uE(n]
Therefore, a' E arg aigmin (aa, x) and the number of elements such that a/z; = E(a', x)) = E(aaaa x) is at most s 1, which contradicts the fact that s is minimal given the RAS choice and Lemma C.3.
As conclusion we have agTk & (a: + 1)z4.
Remark D.2. Recall that Lemma D.1 guarantees that k; defined in (15) satisfy: k; E (1,2) for each i E [n]]
D.1. Proof of Theorem 6.1
Below restate the theorem.
Theorem 6.1. Suppose that Assumption 3.1 holds. Let arg min {(a). in case of multiple optimal actions, consider the one output by RAS when fed with μ. Then, the expected regret of ATA with inputs (B,a) a) satisfies
"
(a ks) (Bp: ((a)) In K,
R& S 2n max BBn - [(a)) + (E(m)
((ā: ki) {(a)))
where a max4e[m] || milley and is a constant.
Proof. Let Ki, be the number of rounds where arm i was queried prior round ls (we take K2 0). Recall that chose the following confidence bound: if K2 2 then
l2k2)2
In(2k?)
conf(i,k) k)) 1
2a
Kip
Ki,
and conf(i, k) = xo otherwise. Recall that A denotes the empirical mean of samples from observed prior to Ki,s 2 and Ar, = if Kg = 0. Let s₁ denote the lower confidence bound used the algorithm:
8i,k (fi,k conf(i, k))
23

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
We introduce the events Es for i [x] and k E [K] defined by
Ei (p pil > conf(i,A))
Let
& = U
Let us prove that for each & E [K] and E [n]: IP (E1,)) S gz, which gives using a union bound P(ε)) S f Let i E [n]. using Lemma E.I and taking 8 = 1/k², we have
P(EL))= = PP|PLK - 14 > conf(i, S k
We call "bad round", a round & where we have E(ax) > {(a). Let us upper bound the number bad rounds.
Observe that bad round there is necessarily an arm [K] such that 0,,K > {(a). For each E [nx], let (a) denote the number of rounds q E (l,..., k} where 0g {()) and arg maxxeln] "j.glf (this corresponds to bad round triggered by worker i) (k) 1gg- E {I....,k) k) aig f(a) and 0i {(a,))
We show that in the case of f(ax)= > {(a). the following event will hold: there exists E [n] such that
24a (ā. + R:)) (n
Bi, E (k
((a, + k:) {(ā))²
To prove this contradiction argument. Suppose that for each E [ma]. - holds and that l(ax) ((ā). This means that & is a bad round, let be an arm that triggered this bad round (i.e., i arg maxje(n) aj,kHly). Event -E, gives in particular
kk)
(1 >
(16)
((a, + k:)p; L(a))
Observe that in each round where N₂(-) is incremented, the number of samples received from the distribution v4 increases by least a + kg. Therefore, we have (16) implies
24a + In(23))
24a In(2k²)
K
((ā, + k:)ps ((a)))
((a) 14 a.
Then we have, using the expressions of conf(.)) and the bound above
1n(222)
1n(2)
2conf(i, = /
+ 4a
Ki
{(a)
$ H4 ā; + k.
(17)
The contradiction hypothesis gives that dia, > ((a), then we have, using the definition of k; that 0i 2 03 his Therefore,
(17) gives
((a)
2conf(i, k) #4
(18)
0g.k
Observe that in each round ||aæ|oo = B, therefore if we have a 2 ā. + K; > à; for some i, we necessarily have that there exists j E [n] \ {i) such that dj.k 0,- 1. Using the fact that {(a) ā,44 with (18), we get
a2 2conf(i, k)) > 0z145
(19)
Since both -E,s and EE hold (because -E,,s implies -Ex), we have that
#4 2conf(i, k) < Úi. conf(i, k) < Si,k1
(20)
24

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
and My u3 conf(jj k). Recall that Myy 0, therefore
Pj 2 (Az, conf(j, = $j.k
(21)
Using the bounds (20) and (21) in (19), we have
01,4 > 0,SJA 2 [aj,k + 1)sg,k
where we used the definition of j in the second inequality. This contradicts the statement of Lemma D.I, which concludes the contradiction argument. Therefore, the event that a bad round implies that E holds for at least one E [n]]. We say that bad round was triggered by arm i, round where (.).) was incremented. Observe that &o E [K] not a bad round then E [e(as)] M() 0, otherwise if is bad round triggered by E [n] then IE [e(ax)] {(a) s Bug {(a). To ease notation introduce for E [nd]
H: : 2a, + k.) ln(2
((ā, + k;)p {()))
The expected regret satisfies
R
B(((ar) -()))
(Bj4 - ((a))
K
Bi4 (( [1 (A: is a bad round triggered by i)]
k=
a
< max sE[n]
(max{(Bpp (Bps {(a))) Σ F(Ex)
Σ (Bp: {(a))
Σ
E [1: (: is bad round triggered by i) -&4]
K
K
I max (aax - {(a))} vE[n]
Σ
P(EA) +
Σ (Bp--E(aj - -((a))
Σ
E[(( = 1+ N. and N $ H,) -E4]
1=1
i=1
K
< max {(ā))}
Σ
P(Ex) + Σ
(Bp - E(ā))
xE(n]
k=1
i=1
< max (Bpu {(a))) 24a2(a. + (B4 ((a)) ln(2
(E[n)
((a: + ki)au - {(a))²
D.2. Proof of Theorem 6.3
Theorem 6.3. Suppose that Assumption 3.1 holds. Let ā E arg mi E(a), in case of multiple optimal actions, we consider the one output by RAS when fed with μ. Then, the expected regret of ATA-Empirical with the empirical confidence bounds using the inputs (B, m)) satisfies
p2 + -B)u InK,
R & 2n max (Bpp-(ppaaj) {(a)} c. Σ VE(n)
((ā: ki) {(a))²
where 17 = maX4e(n) as/ps, and e is constant.
Proof. We build on the techniques used in the proof of Theorem 6.1. Recall the expression of ry:
max 0g
sE[n]
14
Define the quantities C by
In(2k2)
ln(222)
C2 2/
Ki,
Ki,
25

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
Recall that the lower confidence bounds used here defined
$1,6 = -i,k (1 11Cx,h4+
We additionally define the following quantities
(1 Di. (1+ 4 F(Cina
We introduce the events E for i E [n] and & E [K] defined by
{uCuR S and 14 > unk
&,p: {m4- < 81,8
Let
& UselnjE
We have, using Lemma E.3, for each [K] and i E [n]: P (E S g2 which gives using union bound P(Ex) I ka-
Following similar steps in the proof of Theorem 6.1, we call "bad round", round l where we have E(ax) > E(a). Let us upper bound the number of bad rounds.
Observe that bad round there is necessarily an arm i E [K] such that 0,8/4 > ((a). For each E [nx]. let NN(k) denote the number of rounds E [k) where ar,9 > {(a) and E arg max {A (this corresponds to a bad round triggered by worker q):
(k) 1g E (1....,k)) 0igl
((a) and 0i,e4: ((a,))]
We show that in the case of f(ax)= > ((a). the following event will hold: there exists E [n] such that
E, == EN N..(k 1) 185n/ pz k:) In(2 In)
((a: + k:)j4 - {(a)))
To prove this, suppose for a contradiction argument that we have for some E [m]] -Fi and that l a bad round triggered by arm (i.e., {ass) > {(a) and E arg max (gxMz))
This gives in particular
185 12 + k2) ln(2k
N.(k - 1)>
(22)
((a: + k:) - {(a))
Observe that in each round where N.(-) is incremented, the number of samples received from the distribution increases by at least a. + kg Therefore, we have (22) implies
185 185m2 (.)2
185m72 (n(2
K2,
((a: + k2) - {(a))²
((a) o,tk
Therefore, we have, using the expression of C₁ and the bound above
12(2)2)
In (2k22
Ci = 2
+ 22
Kid
Kip
{(a)
< 191j
#4
(23)
ā; + ki
The last bound implies in particular that 1/C, S
3 hence (1 C)Ctt) n
I
26

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
We have
Pic (1 -1 (()o,)))))
Pi.k 81,6
S Nelaiiil
n
14i
1
((a)
$ 3
Ni
(24)
+ Kis
where we used the event -E, in the penultimate inequality (in particular š1,5 = fit,x n(C,,)+ S μe). and the bound (23) in the last inequality. Given the hypothesis that {(ax) >> ((ā)) and, {(ax) 0j.k/4s, we necessarily have 0i,k ā + kj Therefore, bound (24)
{(a)
5 - 5 8i,k < 14: 0g,b
Observe that each round ||axloo = B, therefore we have 04 2 04 + K2 > 0g for some i, we necessarily have that there
exists j E [n] \ {i} such that aj.k S aj - 1. Therefore, using the fact that ((a) 2 aiHl,, we obtain
ag (1s + 58 5)(,)) {(a) 2 C,₁y
(25)
Since both -E,k and E5,e hold (because -Fi,e implies -Ex,, we have that
M4 + 58 5p
81,k + #4 -ii + 4 ( - (si,))
81,4 + 14 -ii + 4ftc ((1 010 -
< 8i,k + 14 -ii, -4ir,nnCikk
S 8 + -ijkk
To conclude, observe that given nCC
3 19.
event -ii, implies that μ4 S 0g therefore
14 58, 5 S Bik
Since -Ei, holds, we also have
M3 2 8j.
Using the two last bounds in (25), we have
a1,4 ajoj. 2 (aj, 1)83,h.
where used the definition of j, as arm satisfying aj 1 ajkk in the second incquality. This contradicts the statement of Lemma D.1, which concludes the contradiction argument. Therefore, the event that & is a bad round implies that E holds for at least one i E [n]] We say that bad round was triggered by arm é, round where N,(-) was incremented. Observe that if lo E [K] is not bad round then E [((ax)] f(ā) = O, otherwise if le is bad round triggered by i E [x] then
[(((x))] E(a) < Bp: ((a). To ease notation we introduce for E [n]
H, 185m²2u (a, k. In (2K2)
((a: k:)p: l()))
27

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
The expected regret satisfies
RK=
Bl((as) {(a))
- ( ((0))
E=1
Bj4 ((a)) [11 (A: is bad round triggered by i)]
1=1
K
K
S max iE[n]
(B)x {()))}
Σ
P(Ex)
Σ
BBp -(())
Σ
E [1: a bad round triggered by i) -Ex
1=1
1=1
k=1
K
K
< m - {(a))}
P(Ex) +
Σ
(Bpt - -((a))
Σ
E [[[(N(( = 1 (k and N; S H,) -Ep]
vE[n|
1=1
i=1
[ax {(ā))}
P(Ex) Σ (Bp: ((ā))
xE(m]
185n 2 1r (a. + kB) {(a) ln(K3
s 2n ma
(E(n)
((as + ki) E(a))²
D.3. Proof of Theorem 4.2
Let us first restate the thcorem.
Theorem 4.2. Suppose Assumption 3.1 holds and let 17 maxqe(]] au/ps, where a =||x; Milly, Then, the total expected computation time after K. rounds, using the allocation prescribed by ATA with inputs (B, a) satisfies
Cx $ (1 + (n Olnn)).
Proof. Let Eg be the expectation with respect to the variables observed up to and including & and Fs the corresponding filtration. Using the tower rule, we have
K
Σ [[((a))] = E
Σ E [((x))]
Consider round [K], let us upper bound Eg- (C(a)) using E, [((a))]. Recall that a E Fx-1, let let Een x(u) since Y, is the sum of a i..d samples we have that 1 = a and || an,pHllle I ai ||X, pellar Thus, using Lemma E.4, we get
asab
Ex-1 [C(az)] Ex-1
max (Esupp(a))
(a)
max (Esupp(as)
{aiktlit + 4 max (Esupp(as)
{ar,a,} II(B)
S max
{ai,ulii +4 max (Esupp(as)
{ar, m4} In(B)
(Es))p(a.))
(1 + An) max {agam))
Moreover, using Jensen's inequality, we have
|
max
max iE
Σ x (w)
EEC(a*)
VE Tq
Using the last two bounds with the result of Theorem 6.1, we get the result.
28

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
E. Technical Lemmas
The lemma below gives a concentration bound on sub-exponential variables. Note that this result can be inferred from Proposition Maurer Pontil (2021), although applying the last result directly requires assuming the variables are positive, this needed in their proof in the one dimensional case. For completeness, we present the full proof below.
Lemma E.1. Let Yj..... Y, be ii.d random variables with IE [Vi] = 0 and a \ p +00. Then for any E (0, 1). have with probability least 1
(n(2/6) n
1n(2/5)
Y
2a
2n
Proof. Let == 2na2. We have using Lemma E.5 that
2na and Σ E
[00
- o!
vay-2 2
Therefore, using Bernstein concentration inequality (Proposition E.2) we obtain that
P
Σ
Y 2 2a + at < 2exp((-t)
Choosing = ln(2/6), we obtain the result.
Proposition E.2 (Theorem 2010 in Boucheron et al. (2005)). Let XI,..., Assume there exist positive numbers and such that
Xn be independent real-valued random variables.
and Σ [{(()))] S g! vc
Σ E E
S
for all integers q 2 3,
where x4 max 0}. Define the centered sum
: Σ (X; EX,) 1=1
Then, for every > 0,
PP (S2 /2 + +el) se
Lemma E.3. Let XI,...;X, Xn be i.i.d positive random variables with mean l and [X +00. Denote
[| e Denote Xn, Et X and let d E 10. Define C.,, by:
: a
Cn, = 2/ (n(2/)) and
+ In(2/8) n
n
Then with probability at least 1 - 8 we have
8n Cnn.))+
where the notation (a)+ max(0, Moreover, in Cn,5 S 1.- then with probability least 8
4
- S 8, 1+ 5" Cn,6
Proof. Fix n, d. We work on the event
En,b = {|xn -1 1 S Cn.8
29

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
that holds with probability least 1 if we apply Proposition E.1 X; μ.. Proof of Xn, (1 - 17 Cn,6)+:
If n,,5 2 1, we have (1 1 Cn,5)+ and the result follows from the fact that X is non-negative which gives μ 2 0. Suppose now that ryCn,5 < 1. Recall that event εn, implies that
Xa S 11l + nyCn.6).
Therefore, have
Xn 1 C/Cn p
Using 1 - nnn8 T+mOn.
with the bound above, we obtain
&n
8n -nn,s++ -Xm(1-)Cnn) S
S
1 ny. Cv.6
Proof of S (1+ 3 Cn.6)
Recall that event En, gives
Xn M M aOn,s u(l mnn,)))
Suppose that n/n, S 4- We therefore have
1.
Xn
1/Cn,5
Next, we use the fact that for any x E [0, 1/4). we have
1
4 3 g
S
1
+
1 -
which gives
4 1 ( Cn.5) 1 S 8n
when ryCn, S 4'
Lemma E.4. Let XI,..., Xa be sequence of independent random variables with finite Orlicz norm [Xil < to0 and let E [X.] = py- Then we have
E[max X.] < max pa + 4a In(n),
Te(m
sE|n]
where a max ||
14 veller
Proof. If n the bound is straightforward, suppose that 2 2. Let Y, X; ps. then max (illly Let A (0, 1/a), we have
max Y. = sE|n]
exp
A max Yi sE(n)
Σ exp (AYs)
Taking the expectation and using Lemma E.5, we have
E
max
In
(exp (AY;)]
sE(n]
1 In
Xae
30

---

ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning
We choose À 1-1/n which gives
max Y [E(n]
In(n²)
1 1
= 2a
n
In(n)
1
< 4a In(n).
(26)
Let E arg maX]e]]] X;, we have
max EEIn]
X, max = X². max
iE|n]
sE(n]
S X. 14+ I max X, 14: = max Y
iE[n]
iE(n]
Combining the last bound with (26) we obtain
max X₃
max 44 + 4a (nnl..)...
iE[n]
iE(n
Lemma below is based on standard argument we give here for completeness.
Lemma E.S. Let Y be a variable such that a = ||Y| < +00. Then we have for any X E (-2 m
E (exp (AY)) s 1 aa/a
Moreover, we have for any q 2 3
E [ - 9!
2
< 2a²
and
2 (2a2) a
Proof. Let 2 = /a) First observe that we have
h
Σ 220
(exp(2)
E [exp (Yxp((Y/aa) 22,
k!
s0,
E k
< 1.
k!
This implies E [x]]]]]]] < * for all & 2 1. Using this bound, we estimate the moment generating function
E [exp(AX)] s IE (exp (l>) [Y]]
|^I* E
k!
120
$ 11
| a
1
1
[|] a
| < M
The remaining bounds follow from
31


Reviews:
[{"review_id": "On98HLkOqs", "reviewer": "Reviewer_xCoQ", "summary": "This paper introduces ATA (Adaptive Task Allocation), a method that dynamically adapts to heterogeneous and random worker computation times without prior knowledge of their distributions. Theoretical analysis shows that ATA achieves optimal task allocation, performing as well as methods that have prior knowledge of computation times. Experimental results confirm that ATA significantly reduces resource costs compared to greedy methods, which can become arbitrarily expensive as the number of workers increases.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their valuable feedback.\n\n> missing ablation study (e.g., impact of data-dependent concentration inequality)\n\nWe are not sure what the reviewer means here: We did study the impact of the data-dependent concentration inequality. Indeed, we test both ATA (the algorithm without the adaptive concentration) and ATA-Empirical (the one with the adaptive one).\n\n> realistic simulation (when workers have heterogeneous and time-vayring computation speeds)\n\nWe did additional experiments with heterogeneous time distributions for the workers, using five types: Exponential, Uniform, Half-Normal, Lognormal, and Gamma. \nThe plots are available\n[here](https://anonymous.4open.science/api/repo/hetero_plots-37B3/file/hetero_plots.pdf?v=e4d1e678).\nThe results show that the algorithms are robust across different distributions.\n\t\nAlso, the computation speeds are time-varying because they are stochastic. It is true that we did not consider the setting of distribution changing over time, however, we believe that, given this is -to the best of our knowledge- the first contribution addressing adaptive allocation with partial feedback and theoretical guarantees, it is both natural and appropriate to begin with the stationary setting.\n\n> Is there any case that information of previous runs degrade performance or slow down the convergence? Could you elaborate more on the impact of using prior information?\n\nOn page 8, we discussed the possibility of using prior estimates of computation times from previous runs.\nWe add more details here on this possibility:\nIf the prior estimates are completely off, the algorithm could make wrong choices in the first rounds, however sooner or later the exploration (that is always active) will cause the algorithm to gather enough data to converge to the correct estimate of the means. Hence, the asymptotic guarantee will be the same and prior knowledge almost always accelerates convergence.\nTo validate this, we conducted additional experiments, please check the plots \n[here](https://anonymous.4open.science/api/repo/prior_knowledge-2E5C/file/prior_knowledge.pdf?v=b620dfe3).", "claims_and_evidence": "The general task allocation problem in parallel stochastic optimization leads to resource wastefulness, and addressing it can improve efficiency across various distributed machine learning methods. Supported by many asynchronous SGD methods such as Ringmaster ASGD and Rennala SGD.", "methods_and_evaluation_criteria": "**Methods** This work formalizes the task allocation problem as a combinatorial online learning problem with partial feedback and non-linear losses. The authors propose ATA, a lower-confidence bound-based algorithm that is agnostic to workers' computation times and achieves near-optimal computation efficiency. They also introduce ATA-Empirical, a data-driven variant that improves empirical performance, and validate their approach through numerical simulations.\n**Evaluation** Authors test the algorithm by simulating a scenario with n workers. They collect 23 gradients from the workers and perform gradient descent step. They compare with GTA, and UTA.", "theoretical_claims": "Proposed algorithms and theorems are backed by theorietical proof.", "experimental_designs_or_analyses": "For different number of workers, they measured runtime vs suboptimality, suboptimality against total worker time, average iteration time and averaged cumulative regret.", "relation_to_broader_scientific_literature": "It mainly focus on asynchronous training field.", "essential_references_not_discussed": "Most of recent works are all referenced.", "other_strengths_and_weaknesses": "(+) solving this problem in online bandit problem is interesting. \n(+) strong theoretical support on the proposed algorithm\n(-) missing ablation study (e.g., impact of data-dependent concentration inequality), realistic simulation (when workers have heterogeneous and time-vayring computation speeds)", "other_comments_or_suggestions": "Please see above", "questions_for_authors": "Is there any case that information of previous runs degrade performance or slow down the convergence?\nCould you elaborate more on the impact of using prior information?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jYGfnw3XiH", "reviewer": "Reviewer_9eUV", "summary": "This paper suggests an adaptive method for task allocation, Adaptive Task Allocation (ATA). In the context of machine learning, this approach addresses federated learning, where in each round, $n$ workers collaboratively perform minibatch SGD with total batch size of $B$. In contrast to the standard approach, where each worker processes a fixed minibatch of size $B/n$, the proposed approach aims to provide a more efficient allocation of minibatch sizes across workers that minimizes the overall runtime per round without unnecessary computational overhead. The paper provides theoretical guarantees in the framework of online convex bandit, and presents empirical results to validate the effectiveness of the proposed approach.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review\n\n> The paper provides theoretical guarantees in the framework of online convex bandit\n\nSince the action set is discrete, the proposed reduction does not fall under the framework of online convex bandits but rather corresponds to a combinatorial bandit problem.\n\n> however, given that this is for a machine learning publication, a thorough evaluation (and adaptation) using standard benchmark datasets and more complex objective functions would be more appropriate.\n\t\t\nPlease take a look at our response to reviewer bw5d. We also provide additional experiments\n[here](https://anonymous.4open.science/api/repo/real_data-3BF5/file/real_data.pdf?v=33782a9d).\n\n> Weakness 1\n\t\t\nThis choice does not affect the theoretical results.\nMore in detail, in the initial phase of the procedure, due to the definition $s_{i,k} = (\\hat{\\mu}\\_{i,k} - \\text{conf}(i,k))\\_{+}$, all workers satisfy  $s\\_ {i, k} = 0$. After this warm-up stage, the scores $s_{i,k}$ become strictly positive. Recall that at round $k$, $\\hat{\\mu}\\_{i,k}$ denotes the empirical mean of the computation times for worker $i$, and $\\text{conf}(i,k)$ is a confidence term of order $\\alpha^2 \\sqrt{\\frac{\\ln(2k^2)}{K{i,k}}}$, where $K_{i,k}$ is the number of times worker $i$ has been queried up to round $k$. This implies that the number of rounds for which $s_{i,k} = 0$ is at most $\\alpha^2 \\ln(k) / \\mu_i^2$.\n\t\t\nFurthermore, in the theoretical guarantees--specifically in the upper bound of Theorem 6.1 (see also Equation (11))--the number of rounds in which worker $i$ incurs a loss due to overallocation is of order $\\alpha^2 \\ln(k)/\\left(\\mu_i - \\frac{\\ell(\\bar{a}, \\mu)}{\\bar{a}i + k_i}\\right)^2$. This shows that the losses incurred during the warm-up stage, where $s_{i,k} = 0$, are dominated by those arising when $s_{i,k} > 0$ for all $i$. A more detailed analysis is provided in the proof of Theorem 6.1 (Section D.1).\n\t\t\nIn practice, during the warm-up phase when many $s_{i,k} = 0$, the allocation $a_k$ (line 5 of Algorithm 1) is not uniquely defined. As noted by the reviewer, we resolve this by uniformly distributing the workload among the workers with $s_{i,k} = 0$. However, any heuristic that ensures $a_k \\in \\arg\\min_a \\ell(a, s_k)$ could be used without any change to the theoretical guarantees. We chose the uniform strategy for its simplicity, and as our experiments show, it applies only during a small fraction of the early rounds. We will add this clarification as a remark in the main body of the paper.\n\t\t\n\t\t\n> Weakness 2\n\t\t\nThis is a limitation of virtually any bandit algorithm. In fact, in the case of bounded distributions, many existing procedures assume knowledge of the range of the losses. Moreover, we do not believe it is possible to design a completely parameter-free procedure for the problem we considered because proving any concentration result on the empirical mean of a random variable without any assumption on its variance is known to be impossible.\n\nThat said, we would like to draw the reviewer's attention to the fact that ATA-Empirical requires only $\\eta$ as input, instead of $\\alpha$. That is, each of our algorithms requires only a single distribution-dependent parameter. For example, in cases where the learner knows that the computation times approximately follow an exponential distribution--which is a common modeling assumption (see lines 152–160)--it is advantageous to use ATA-Empirical, since $\\eta$ is necessarily of constant order (i.e., $O(1)$).\n\t\t\nRegarding the warm-up phase where $s_{i,k} = 0$: as with ATA, in ATA-Empirical, a larger value of $\\eta$ results in a longer initial phase during which $s_{i,k} = 0$. However, the previous observation still holds: The loss incurred during this stage is asymptotically dominated by the regret bound guaranteed by the algorithm.\n\t\t \n> Weakness 3\n\t\t \nWe are not entirely certain we understand the reviewer’s point. The confidence bounds used for the computation times *increase* with the number of iterations (ensuring exploration) and *naturally shrink* with the number of observed samples from the corresponding distribution (this quantity is denoted $K_{i,k}$ for worker $i$ up to round $k$). These bounds are derived using concentration inequalities for sub-exponential distributions (see Proposition E.1).\n\t\t\nAlso, our theoretical guarantees are established under the assumption of stationary distributions, hence, the expected value of the workers' running times cannot change. We believe that, given this is--to the best of our knowledge--the first contribution addressing adaptive allocation with partial feedback and theoretical guarantees, it is both natural and appropriate to begin with the stationary setting. This assumption is also standard in the majority of bandit literature. We leave extending the analysis to the non-stationary case, where workers' distributions may vary over time, as an important direction for future work.", "claims_and_evidence": "The proofs of the theoretical claims were not checked, and the empirical results may be less convincing in realistic scenarios (see \"Methods And Evaluation Criteria\" and \"Weaknesses\").", "methods_and_evaluation_criteria": "The authors evaluated their approach using regret analysis, which is standard for bandit applications. They also provided an empirical evaluation in a simple convex optimization setting; however, given that this is for a machine learning publication, a thorough evaluation (and adaptation) using standard benchmark datasets and more complex objective functions would be more appropriate.", "theoretical_claims": "I didn't check the correctness of the proofs; the complete proofs are in the appendix and weren't reviewed.", "experimental_designs_or_analyses": "I reviewed the experimental design and analysis in the paper (issues were discussed at \"Methods And Evaluation Criteria\" and \"Weaknesses\").", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "*Strengths*:\n1. I found it interesting to approach this problem from the perspective of a bandit problem, which, on one hand, exploits fast workers and, on the other hand, allows for exploring changes in the performance of weaker workers, whose running time can vary over time. This can be realistic in some practical scenarios. The proposed approach might be a good first direction, but modifications in the confidence assessment may be necessary for realistic and more complex problems (see weaknesses).\n\n2. This idea can be relevant today in the context of energy efficiency, as GPUs are used at scale and consume significant amounts of power.\n\n*Weaknesses*:\n1. In the case where $s_{i,k} = 0$ for some worker $i$,  the authors used a suboptimal uniform allocation (see \"Other Comments Or Suggestions\") instead of an impractical optimal allocation (e.g., assigning $B$ allocations to some worker $i$ with $ s_{i,k}= 0$ and $0$ to the others). How does this choice affect the theoretical results?\n\n2. ATA-Empirical relies on the parameter $\\eta$, which is unknown in practice and needs to be evaluated. This can be critical, as large $\\eta$ may cause $s_{i,k}$ to become zero for some worker $i\\in[n]$, which leads to a suboptimal uniform allocation.\n\n3. The confidence assessment depends on the number of iterations $k$, $K$. While using these parameters might work for small and simple problems, as demonstrated in the experiments, it may be less practical in more complex scenarios. In such cases, the confidence assessment may not align with the actual changes in workers' running times, and make it less meaningful.", "other_comments_or_suggestions": "* The case where $s_{i,k} = 0$ for some worker $i$ is only discussed in the appendix on page 16, lines 841–842, instead of at the main paper.\n\n* To avoid being misleading, the authors should explicitly include the dependence of Compute LCBs on $\\alpha$ or $\\eta$ in the proposed algorithms (such as in Algorithm 1).", "questions_for_authors": "see weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LYBZymDGFy", "reviewer": "Reviewer_svBT", "summary": "This paper proposes a new task allocation algorithm for efficiently managing computation resources in distributed SGD.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their valuable feedback.\n\n> The paper only shows the simulation results, but there is no real dataset experiment (e.g. distributed deep learning training allocation) to show the performance of their approach.\n\t\t\nThere might be a misunderstanding on the nature of our contribution: We propose an allocation strategy across workers that adapts to their computation times and that is *agnostic* to the specifics of the task.\nThis means that the behaviour of our algorithm does not depend on the optimization algorithm used, the dataset, the loss function, etc.\nInstead, it *only* depends on the observed computation times. This means that, given a sequence of computation times for any number of workers and a training loss curve, anyone can immediately simulate how our algorithm would affect the total computation time. Indeed, this would even work for other tasks different than optimization.\n\nThat said, to answer to the reviewer's point, we ran additional experiments using a CNN trained with Adam on the CIFAR-100 dataset.\nThe plots can be found\n[here](https://anonymous.4open.science/api/repo/real_data-3BF5/file/real_data.pdf?v=33782a9d).\n\n> Given that the computation time of worker nodes may be affected by dynamic load, task switching, etc. The distribution assumption and bound of the nodes may not be practical in real-world scenario.\n\t\t\nWe acknowledge the reviewer's point that incorporating dynamic load and task switching would align the model more closely with practical scenarios. However, to the best of our knowledge, this is the first work on adaptive online allocation under a limited budget that provides strong theoretical guarantees on the optimality of the allocation strategy. It seems a bit unfair to require the solution of all the possible problems in a single first paper.\nThe assumptions made regarding the distribution of computing times are standard and commonly adopted in the literature for modeling such quantities. We leave extending the framework to accommodate more general and irregular distributions for future work.\n\n> It seems that there is no computational complexity analysis of the proposed method, and the comparison with baselines.\n\nWe do discuss the computational complexity, but maybe we could stress it more. Specifically, for our main procedures, ATA and ATA-Empirical, the computational complexity per round is primarily determined by the step of computing the allocation $a_k$ (line 5 in Algorithm 1). As stated in lines 205-208, the complexity of this step is linear in $n$ when the number of workers $n$ exceeds the task budget $B$, demonstrating the efficiency of our approach. For more details please check Section C of the appendix.\n\n> The experiment section lacks explanation of parameter settings.\n\nThank you for pointing this out. We noticed that the dimensionality and stepsize were missing. For our experiments, we chose $d=100$ and a stepsize of $\\gamma=10^{-4}$. Is there anything else we may have missed? We will include all the missing information in the camera-ready version of the paper.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "I have read the theoretical claims in the main part. The proofs in the appendix are not checked.", "experimental_designs_or_analyses": "Authors implemented simulations to evaluate the validity of the proposed algorithm. But there is no real dataset experiment.\n\nThe paper lacks explanation of parameter settings in section 7 (experiment part).", "relation_to_broader_scientific_literature": "The paper reduces task allocation problems to an online bandit problem (MAB problem), and utilizes proxy loss to solve the computational challenge caused by the combinatorial nature of the action set.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strength: The proposed algorithm has solid theoretical analysis. The authors innovatively reduce the task allocation problem to a MAB problem.\n\nWeaknesses: \n1. The paper only shows the simulation results, but there is no real dataset experiment (e.g. distributed deep learning training allocation) to show the performance of their approach.\n2. Given that the computation time of worker nodes may be affected by dynamic load, task switching, etc. The distribution assumption and bound of the nodes may not be practical in real-world scenario. \n3. It seems that there is no computational complexity analysis of the proposed method, and the comparison with baselines.\n4. The experiment section lacks explanation of parameter settings.", "other_comments_or_suggestions": "Could you add some simple experiments for real dataset? Please explain the parameter settings in section 7.", "questions_for_authors": "No more questions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hkH8Wi9zZm", "reviewer": "Reviewer_bw5d", "summary": "In distributed machine learning, traditional greedy methods waste resources by assigning tasks to every worker, even when only a few results are needed. This paper tells a story of efficiency: it introduces ATA, a smart task allocator that learns the speed of each worker on the fly. Instead of needing prior knowledge, ATA treats task allocation like a multi-armed bandit problem, using lower confidence bounds to decide where to assign tasks. By balancing exploration (learning each worker’s speed) and exploitation (favoring faster workers), ATA nearly matches the performance of an ideal, all-knowing strategy while eliminating wasted work. An improved variant, ATA-Empirical, tailors this approach even further to individual workers, achieving better results in practice. Overall, the paper offers a principled and practical solution to cut waste and boost efficiency in large-scale machine learning.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review\n\n> they lack real-world experiments on real-world benchmarks such as MLPerf.'}\n\nThere might be a misunderstanding on the nature of our contribution: We propose an allocation strategy across workers that adapts to their computation times and that is *agnostic* to the specifics of the task.\nThis means that the behaviour of our algorithm does not depend on the optimization algorithm used, the dataset, the loss function, etc.\nInstead, it *only* depends on the observed computation times. This means that, given a sequence of computation times for any number of workers and a training loss curve, anyone can immediately simulate how our algorithm would affect the total computation time. Indeed, this would even work for other tasks different than optimization.\nBesides, MLPerf does not contain benchmarks for distributed optimization or bandit algorithms.\n\nThat said, to answer the reviewer's point, we ran additional experiments using a CNN trained with Adam on the CIFAR-100 dataset.\nThe plots can be found\n[here](https://anonymous.4open.science/api/repo/real_data-3BF5/file/real_data.pdf?v=33782a9d).", "claims_and_evidence": "The paper claims that ATA dynamically learns worker speeds and allocates tasks so as to minimize total computation time without prior knowledge, nearly matching the performance of an ideal oracle. It recasts task allocation as a combinatorial bandit problem with rigorous theoretical guarantees, ensuring logarithmic regret and near-optimal performance, and shows through experiments that both ATA and its variant, ATA-Empirical, substantially reduce wasted computation compared to greedy methods.", "methods_and_evaluation_criteria": "Evaluation is carried out through simulations on distributed SGD tasks, comparing the proposed methods against baselines such as Greedy Task Allocation, Uniform Task Allocation, and an oracle-based fixed strategy. Key metrics include runtime efficiency, total worker time, average iteration time, and cumulative regret.", "theoretical_claims": "n/a", "experimental_designs_or_analyses": "The experiments of Figure 1 are conducted on a variety of algorithms and use a wide range of metrics to measure the effectiveness of these algorithms. While I find these results interesting, they lack real-world experiments on real-world benchmarks such as MLPerf.", "relation_to_broader_scientific_literature": "The paper builds on asynchronous optimization methods like Rennala SGD and Ringmaster ASGD by addressing their resource inefficiencies through a novel formulation of task allocation as a combinatorial multi-armed bandit problem. It leverages foundational bandit theory, particularly lower confidence bounds, to balance exploration and exploitation without prior knowledge of worker speeds, while also resonating with federated learning approaches that adjust computations for heterogeneous devices. This integration of ideas from distributed optimization, bandit theory, and federated learning marks a significant advancement in tackling resource waste in large-scale machine learning.", "essential_references_not_discussed": "n/a", "other_strengths_and_weaknesses": "n/a", "other_comments_or_suggestions": "n/a", "questions_for_authors": "n/a", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes an adaptive method for task allocation, which has an application in collaboratively performing mini-batch SGD with a fixed total batch size. More precisely, unlike standard approaches where each worker processes an equal-sized mini-batch, the proposed adaptive task allocation method dynamically allocates mini-batch sizes across workers to minimize the overall runtime per round, while avoiding unnecessary computational overhead. The underlying technique is a bandit algorithm. The topic is highly relevant to the optimization and machine learning community. All reviewers unanimously recommend acceptance of this work.

The authors are expected to include implementation details for their real-world dataset experiments in the rebuttal phase and incorporate the constructive feedback that they received here in the final version.
Decision: Accept (poster)