Learning Mixture Models with Simultaneous Data Partitioning and Parameter Estimation
Authors: Parth Vipul Sangani
Arjun Shashank Kashettiwar
Durga S
Ganesh Ramakrishnan
Rishabh K Iyer
Abir De
Venue: Submitted to ICLR 2023 (2023)
Keywords: mixture models
resource constrained learning

Abstract: We study a new framework of learning mixture models via data partitioning called PRESTO, wherein we optimize a joint objective function on the model parameters and the partitioning, with each model tailored to perform well on its specific partition. We connect PRESTO to a number of past works in data partitioning, mixture models, and clustering, and show that PRESTO generalizes several loss functions including the k-means and Bregman clustering objective, the Gaussian mixture model objective, mixtures of support vector machines, and mixtures of linear regression. We then propose a new joint discrete-continuous optimization algorithm which achieves a bounded approximation guarantee for any general loss function, thereby achieving guarantees for the afore-mentioned problems as well. We study PRESTO in the context of resource efficient deep learning, where we train smaller resource constrained models on each partition and show that it outperforms existing data partitioning and model pruning/knowledge distillation approaches, which in contrast to PRESTO, require large initial (teacher) models.

Under review as a conference paper at ICLR 2023
LEARNING MIXTURE MODELS WITH SIMULTANEOUS DATA PARTITIONING AND PARAMETER ESTIMATION
Anonymous authors Paper under double-blind review
ABSTRACT
We study new framework of learning mixture models via data partitioning called PRESTO, wherein we optimize a joint objective function on the model parameters and the partitioning, with each model tailored to perform well on its specific parti- tion. In contrast to prior work, we do not assume any generative model for the data. We connect our framework to a number past works in data partitioning, mixture models, and clustering, and show that PRESTO generalizes several loss functions including the k-means, Bregman clustering objective, the Gaussian mixture model objective, mixtures of support vector machines, and mixtures of linear regression. We convert our training problem to a joint parameter estimation cum a subset se- lection problem, subject to matroid span constraint. This allows us to reduce our problem into constrained set function minimization problem, where the under- lying objective monotone and approximately submodular. We then propose a new joint discrete-continnng approximation guarantee for our problem. We show that PRESTO outperforms several alternative methods. Finally, we study PRESTO in the context of resource efficient deep learning, where we train smaller resource constrained models on each partition and show that it outperforms existing data partitioning and model /nowin/g distillation approaches, which in contrast to PRESTO, require large initial (teacher) models.
optimization algorithm which achieves a bounded
INTRODUCTION
In the problem space of learning mixture models, our goal is fit given set of models implicitly
to different clusters of the dataset. Mixture models are ubiquitous approaches for prediction tasks
on heterogencous data (Dasgupta) 1999, Achlioptas & McSherr 2005 Kalai et al. 2010 Belkin
& Sinha 2010a Pace & Barry, 1997: Belkin Sinha [010b] Sanjeev & Kannan 2001 Hopkins
2018 Fu & Robles-Kelly 2008). and find use in plethora of applications, e.g., finance, ge-
nomics (Dias et al. [2099] Liesenfeld 2001: Pan et al. 2003). etc. Existing literature on mixture models predominately focuses on the design of estimation algorithms and the analysis of sample complexity for these problems (Faria & Soromenho 2010 Städler et al. 2010; Kwon et al., 2019; Yi al.) (2044), and analyzes them theoretically for specific and simple models such as Gaussians, linear regression, and SVMs. Additionally, erstwhile approaches operate on realizable settings they assume specific generative models for the cluster membership of the instances. Such an as- sumption can be restrictive, especially when the choice of the underlying generative model differs significantly from the hidden data generative mechanism. Very recently, Pal et al.) (2022) consider a linear regression problem in non-realizable setting, where they do not assume any underlying generative model for the data. However, their algorithm and analysis is tailored towards the linear regression task.
1.1 PRESENT WORK
Responding to the above limitations, we design PRESTO, novel data partitioning based framework for learning mixture models. In contrast to prior work, PRESTO is designed for generic deep learning problems including classification using nonlinear architectures, rather than only linear models (linear regression or SVMs). Moreover, we do not assume any generative model for the data. We summarize contributions follows.

---

Under review as a conference paper at ICLR 2023
Novel framework for training mixture models. At the outset, aim to simultancously partition the instances into different subsets and build a mixture of models across these subsets. Here, each model is tailored to perform well on specific portion of the instance space. Formally, given set of instances and the architectures of K' models, we partition the instances into K disjoint subsets and train a family of K component models these subsets, wherein, each model is assigned to one subset, implictly by the algorithm. Then, we seek to minimize the sum of losses yielded by the models on the respective subsets, jointly with respect the model parameters and the candidate partitions of the underlying instance space.
Note that our proposed optimization method aims to attach each instance one of the K models on which it incurs the least possible error. Such an approach requires that the loss function helps guide the choice of the model for an instance, thus rendering it incompatible for use at inference time. We build an additional classifier to tackle this problem; given an instance, the classifier takes the confidence from each of the K models as input and predicts the model to be assigned to
Design of approximation algorithm. Our training problem involves both continuous and com- binatorial optimization variables. Due to the underlying combinatorial structure, the problem is NP-hard even when all the models are convex. To solve this problem, we first reduce our training problem to a parameter estimation problem in conjunction with a subset selection task, subject to a matroid span constraint (lyer al. 2014). Then, further transform it into a constrained set func- tion minimization problem and show that the underlying objective is a monotone and a-submodular function (El Halabi & Jagelka 2020 Gatmiry & Gomez-Rodriguez) (018) and has bounded cur- vature. Finally, we design PRESTO, an approximation algorithm that solves our training problem. by building upon the majorization-niniiizatii
algorithms proposed in (Iyer & Bilmes 2015 [yer iet al. 2013a Durga et al. (02T). We provide the approximation bounds of PRESTO, even when the learning algorithm provides an imperfect estimate of the trained model. Moreover, it can be used to minimize any a-submodular function subject to a matroid span constraint and therefore, is of independent interest.
Application to resource-constraine, settings. With the advent of deep learning, the complexity of machine learning (ML) models has grown rapidly in the last few years (Liu al.) 2020 Arora
2018 Dar et al. 202 Bubeck & Sellke, 2021 Devlin et al. 2018 Liu et al. 2019; Brown eet al.
ctt al. 2020 The functioning of these models is strongly contingent on the availability of high performance computing infrastructures, e.g., GPUs, large RAM, multicore processors,
The key rationale behind the use of an expensive neural model is to capture the complex nonlinear relationship between the features and the labels across the entire dataset. Our data partitioning framework provides a new paradigm to achieve the same goal, while enabling multiple lightweight models to on low resource device. Specifically, we partition dataset into smaller subslices and train multiple small models on each subslice- since each subslice is intuitively a smaller and simpler data subset, we can train a much simpler model on the subslice thereby significantly reducing the memory requirement. In contrast to approaches such as pruning (Wang el al. 2020a; Lee et al., 2019 Lin et al. 2020 Wang et al.] 2020bb] Jiang et al. 2019 Li et al.] 2020 Lin et al.) (017) and knowledge distillation (Hinton et al.) 2015] Son et al. [0021] we do not need teacher models (high compute models) with the additional benefit that we can also train these models on resource constrained devices.
Empirical evaluations. Our experiments reveal several insights, summarized as follows. (1) PRESTO yields significant accuracy boost over several baselines. (2) PRESTO is able to trade-off accuracy and memory consumed during training more effectively than several competitors, egg. pruning and knowledge distillation approaches. At the benefit of significantly lower memory usage, the performance of our framework is comparable to existing pruning and knowledge distillation approaches and much better than existing partitioning approaches and mixture models.
1.2 RELATED WORK
Mixture Models and Clustering: Mixture Models (Dempster al. 1977 Jordan & Jacobs 1994)
and k-means Clustering (MacQueen] 1967 Lloyd 1982) are two classical ML approaches, and have seen significant research investment over the years. Furthermore, the two problems are closely con- nected and the algorithms for both, i.e., the k-means algorithm and the Expectation Maximization algorithm for mixture models closely related the EM algorithm is often called soft-clustering.
2

---

Under review as a conference paper at ICLR 2023
wherein one assigns probabilities to each cluster. Mixture models have been studied for number of problems including Gaussian Mixture Models (Xu & Jordan 1996), Mixtures of SVMs (Col- lobert et al., 2001] Fu & Robles- Kelly 2008), and linear regression (Faria & Soromenho, 2010 Städler et al. 2010, Kwon et al. 2019 Pal et al. 2022). As we will show in this work, our loss based data partitioning objective generalizer the objectives of several tasks, including k-means clus-
tering (Looyd] [182). Clustering with Bregman divergences (Banerjee et al. (005), and mixture
models with SVMs and linear regression (Collobert et al. 2001 Pal et al. 2022).
Pal
Following initial analysis of the EM for Gaussian Models (Joraan & Jacobs 1994] Xu & Jordan,
1996). mixture models have also been studied for SVMs (Collobert et al. [200]] wherein the authors extend the mixture-mode formulation to take the SVM loss instead of the Gaussian distribution. A lot of recent work has studied mixtures of linear regression models. Most prior work involving mixture models pertaining to regression is in the realizable setting, with the exception of (Pal anpectaaioo-maaxiissaiooo
et al. 2022). Mixture model papers aim at using the expectation- maximisation (EM) algorithm for parameter estimation. Balakrishnan et al. (2017) study the EM algorithm which is initialised with close estimates to ground truth, concluding that it is able to give reasonable fits to the data. When the sample size is finite, (Balakrishnan et al. 2017) show convergence within ly norm of the true parameters. Yi et al. (2014) provide yet another initialisation procedure which uses eigenspace analysis and works for two bimodal regressions. These works assume that the underlying Gaussian distributions have the same shared covariance. Li & Liang (2018) lift this assumption, proving near-optimal complexity. Pal al. (2022) approach the problem of linear regression under the non- realizable setting by defining prediction and loss in the context of mixtures and formulating a novel version of the alternating maximisation (AM) algorithm.
Resource-constraaned learning, In the pursuit of better performance, most state of the art deep learning models are often over-parameterized. This makes their deployment in the resource con- strained devices nearly impossible. To mitigate the problems, several handcrafted architectures such as SqueezeNets (landola et al. 2016; Gholami et al. (0188), MobileNets (Howard et al. 2017, San- aler et al., 2018) and ShuffleNets (Zhang et al.] 2018]] Ma et al.) 2018) were designed to work in mobile devices. Recently, EfficientNet (Tan & Le 2019) was proposed, that employs neural archi- tecture search. However, these models are designed to work on the entire training set, and leave heavy memory footprint.
Simultaneous model training and subset selection. Data subset selection approaches are pre- dominately static, and most often, do not take into account the model's 's current state (Wei et al.]
2014a by Kirchhoff & Bilmes, 2014 Kaushal et al. 2019. Liu al. 2015 Bairi et al. 2015 Lucic et al. 2017 7 Campbell& & Broderick 2018: [outsidis et al., 2013). Some recent approaches attempt solve the problem of subset selection by performing joint training and subset selection (Mirza- soleiman et al. 2020a by Killamsetty et al. 202 lb a Durga et al. (02II. On the other hand, some other approaches (Mirzasoleiman al., 2020a]] Kill ailsattytty al. 2021a) select subsets that approx- imate the full gradient. Further, some of these approaches (Killamsetty et al. 2021b a) demonstrate improvement in the robustness of the trained model by selecting subsets using auxiliary or valida- tion set. Durga et al. (2021) train the model in such way that the validation set error is controlled. All the aforementioned techniques try to determine (for purposes of training), a single subset of the training set. In this paper we propose PRESTO as a method to simultaneously select multiple subsets and also learn mixture of models (with a model corresponding to each subset), in order to improve robustness of the trained model, even in resource-constrained settings.
2 PROBLEM FORMULATION
In this section, first present the notations and then formally state our problem. Finally, instantiate our problem in different clustering and mixture modeling scenarios.
2.1 NOTATIONS
We have N training instances ((a,, y:) S S N). where æ; E X' the feature vector and y E y is the label of the thh instance. In our work, we set X' = Rd and treat y to be discrete" Given K, we denote hgp ... hag as the component models that are going to be used for the K subsets resulting from partition of X. Here 02 is the trainable parameter vector of ha, These parameters are not
For brevity, we present our analysis for the clussification setup. However, our framework is also applicable, as it is, to the regression
setup.
3

---

Under review as a conference paper at ICLR 2023
shared across different models, i.e., there is overlap between 0₂ and ON for k * k'. In fact, the models hgg and hgg can even have different architectures. Moreover, the representation of the output heg can vary across different settings. For example, given æ E X, sign(ho, (æ)) is predictor of for support vector machines with y E {+1} whereas, for multiclass classification, he, (x) provides distribution over y. To this end, we use Ehhg, (z), y) indicate the underlying loss function for any instance (xr, y). We define [A] ,I,...,AA A] for an integer A.
2.2 PROBLEM STATEMENT
High level objective. Our broad goal is to fit mixture of models on given dataset, without making any assumption concerning the generative process, the instances or features. Given family of model architectures, our aim is to learn to partition the instance space X' into a set of subsets, determine the appropriate model architecture to be assigned to each subset and to subsequently train the appropriate models the respective subsets.
Problem statement. We are given the training instances D. the number of subsets K resulting from partitioning 1) and a set of model architectures .g...... hox- Our goal, then, is to partition the training set D into K subsets SkSk with Sy n Skr = y and UgSk D that when the model he. is trained on Sş for E [K], the total loss minimized. To this end, we define the following regularized loss, i.e.
F([(S, 0s) k [K]])
E MEE
[e (ho, (zx). v)) + A110x113]
(1)
kE[K] SES,
Then, we seek to solve the following constrained optimization problem:
minimize F(((Sk, 0x) & E [K]}) subject to, S5 n Sw = 0, Ak E [K], U = D. (2)
0 Ox.
Sx
Here, is the coefficient of the Le regularizer in Eq. 2 The constraint S₁ n S.r 0 ensures that that each example i E D belongs to exactly one subset and the constraint [ St: = DD entails that the subsets {Sk k E [K] cover the entire dataset. The above optimization problem is joint model parameter estimation and a partitioning problem. If we fix the partition {S1, S2,.,, Sk}. then the
optimal parameter vector Og depends only on Sg, the subset assigned to it. To this end, we denote the optimal value of 6 for the above partition (Sx) and transform the optimization (2 into the following equivalent problem:
minimize F([S,, (5()) l [K]]) subject to, S₁ n S 0, Vk E [K]. USE(K Se D. (3)
St SL.S2....Sx
Note that, computing (Sk) ) has a polynomial time complexity for convex loss functions. However, even for such functions, minimizing F as defined above is NP-hard.
Test time prediction. Since computation of the optimal partitioning requires us to solve the opti- mization 33. it cannot be used assign an instance & to a model ho, during the test time. To get past this blocker, we train an additional multiclass classifier Ts X - [K], which is trained on (ax, k)} pairs where E Sg, that, during test time, it can assign unseen instance æ to model component ha, using & = (z)
2.3 INSTANTIATIONS IN CLUSTERING AND MIXTURE MODELING
In this section, we show how the formulations listed in the previous section (e.g., Eqn 11) have appeared in several applications ranging from clustering to mixture modeling.
K-Means Clustering: Since k-means clustering (Lloyd) 1982
MacQueen 1967 is unsuper- vised, we do not have access to labels y, It then turns out that 0g = Mgg the cluster means, and ((hg, (x:)) = ||æ; 12.1.
Bregman Clustering: K-means clustering is a special case of the more general clustering with Bregman divergences (Banerjee al. 2005). Again, in this case, the parameters 0g A4k, the cluster means and (z,,, B
Mixture of SVMs: The mixture of Support Vector Machines (Fu & Robles-Kelly 2008 special case of Eq. 1 with ((ho, (z.). = max(0, - Ne, (2r))
"The Bregman divergence defined as: Bo(2, v) =<(>) -b() (Vo(v), x v).
4

---

Under review as a conference paper at ICLR 2023
Mixture of Linear Regression: The mixture of linear regression models (Pal et al. 2022 is special case of Equ. ①I with ((ho (at), 91) = l (æ:) - vell?.
2.4 APPLICATION IN THE RESOURCE CONCCER-NNSTTINED SETUP
In general, the relationship between an instance x: and the label y can be arbitrarily nonlinear. A complex deep neural network captures such relationship by looking into the entire dataset. However, such networks require high performance computing infrastructure for training and inference. The training problem (3) can used build K lightweight models ho...... hox- each of which is localized to specific regime of the instance space X. Thus, a model is required to capture the nonlinearity only from the region assigned it and not the entire set X. As result, can be lightweight and used with limited resources.
Consider a large model with number of parameters equal to the total number of parameters collec- tively across the models, ie., Σx= dim(0x). Such model has to be loaded entirely into GPU RAM for training inference, and thus requires larger GPU RAM. In contrast, our approach quires us to load at time, only one model component he, and the corresponding subset Sy during both training and test. This is instead of having to load all the K model components and the entire dataset. While this can increase both training and inference time, it can substantially reduce the memory consumption by 1/K times in comparison large model having similar expressiveness.
3 PRESTO: PROPOSED FRAMEWORK TO SOLVE THE TRAINING PROBLEM 3
In this section, we first show that the optimization problem (3) is equivalent to minimizing mono- tone set function subject to matroid span constraint (Iyer et al. 2014 Schrijver et al. 2003). Subse- quently, we show that this set function is a-submodular and admits a bounded curvature. Finally, we use these results to design PRESTO, an approximation algorithm to solve the underlying constrained set function optimization problem. We next present these analyses, beginning with the necessary definitions about monotonicity, a-submodularity and different matroid related properties.
Definition 1. (1) Monotonicity, a-submodularity and generalized curvature: Given ground set V, let 2 + R be set function whose marginal gain is denoted by G(e S) G(Su(())
G(S). The function G is monotone non-decreasing if G(e ()) > for all S C and V.S.
G is a-submodular with a E (0, 1] if G(e S) 2 aG(eITT for all S C T and e E V\T (Hashemi et al. 2019: EII Halabi & Jegelka 2020 The generalized curvature of G(S) is defined as (G(
1 minees Gle [S)e)/G((e 0) (yer et al. 2013b Zhang & Vorobeychik 2016). (2) Base, rank and span of matroid: Consider a matroid M (V,2) where Z is the set of independent sets (Refer to Appendix A for more details). A base of .M is a maximal independent set. The rank
function T. 2 N is defined as: (S) = max m [4] A set S is a spanning set it is
the superset of base equivalently, TM (V(VV (Schrijver et al. 2003 Iyer et al. 2014
Edmonds 2003).
3.1 REPRESENTATION OF 3 AS A MATROID SPAN CONSTRAINED SUBSET SELECTION TASK
Transforming partitions into 2D configurations. Given the training instances D and the size of partition K, we first define the ground set V in the space of Cartesian products of D and [K], i.e,, V = D x [K]. Thus V consists of all pairs {(i, k) i E D, k E [K]} which enumer- ates all possible assignments between the instances and model components. Moreover, we define V. kii,k) l E [K]} and V., = ((i, k) E D}. Here, W., {i} x [K] enumerates all pos-
sible assignments of the tth instance and V. = D x {k} enumerates all possible configurations specifically wherein S% is assigned to an instance.
Reformulating optimization 33) in constrained subset selection. Having defined the ground set V and the partitions in 2D configuration space above, define $ {i, k)) Sk, k E [K]} that encodes the set of assignments in space of induced by the underlying partition. Then, the set S. = ((i, k) i e S& } specifies the set of instances attached to the subset & and elucidates the subset containing i. It can be observed that Sy S n V Since every instance is assigned exactly to one subset Sş E {SI...., S J. we have that |S n V = 1. To this end, we introduce the following set function, which is the sum of the fitted loss functions defined in Eq. (I) trained over individual
5

---

Under review as a conference paper at ICLR 2023
subsets.
Σ Σ ( (Mer:env.)) + A|10x|??
G(S))
(4)
kE[K] (SnV).
We rewrite optimization problem as follows:
minimize (SS) subject to, |Sn = Vie
(5)
SCV
Theoretical characterization of the objective and constraints Eq. G). Here, we show that objective admits monotonicity, a-submodularity and bounded curvature in wide variety of scenarios (Proven in Appendix A.2
Given the set function G(S) defined in Eq. (4) and the individual regularized loss functions l introduced Eq. (T). we define: Emin maxk, Eigen VV E(ho, (xx)), 9:)], l = min,g [((h,,( v1) + 1| and Iiin = max ming, (x.)) V4) + Al 0x 11 Then, we have the following results:
Theorem 2.
(1) Monotonicity: The function G(S) defined in Eq. 4 is monotone non-decreasing in S.
(2) e-submodularity: If ((o, (x)), y) is L-Lisschitz for all k E {1,.., K} and the regularizing coef- ficient X satisfies: X Emiin function G(S) is a- a-submodular with
Emin
0G
Emin +
2L2
XL
(6)
(A+0.5mmiu)
(3) Generalized curvature: The generalized curvature (S) for any set S is given by: KG (S) KiG 1 Smiu/Cmnn
Solving the optimization 5 is difficult due to the equality constraint. However, as suggested by Theorem (1). G(S) is monotone in S. Hence, even if we relax the equality constraints [SnV..|
to the incquality constraint V I 2 1, they achieve the equality at the optimal solution S*. Thus, the optimization (5) becomes
minimize G(S) subject to, |S V 2 ViED. SCV
As we formally state in the following proposition, the above constraint (set) can be seen as a matroid span constraint for partition matroid. This would allow us to design approximation algorithm to solve this problem.
Proposition 3. Suppose that set S satisfies [Sn Vial > 1 for all i E D. Then is a spanning set of the partition matroid M (V,I) where I ( 1). Moreover, if S satisfies [S [. for all E D, then S is a base of M.
The above proposition suggests that the optimal solution S* of the optimization 7 is a base of the partition matroid M = (V, I) with {/] n V. S 1}.
3.2 PRESTO: AN APPROXIMATION ALGORITHM SOLVE THE OPTIMIZATION PROBLEM
In this section, we present our approximation algorithm PRESTO for minimizing the optimization problem (5). which is built upon the algorithm proposed by Durga et al. (2021). Their algorithm aims approximately minimize an a-submodular function, whereas extend their algorithm for min- imizing an a-submodular function with matroid span constraint and derive an approximation guar- antee for it. Specifically, we employ the Magjorization-Mnnmmizze for minimizing the optimization problem Toward that goal, we first develop the necessary in- gredients as follows.
approach (Iyer et al.) 2013a b)
Computation of modular upper bound. First, we present modular upper bound for G(S) Iyer
al. 201 Given fixed set 3 and the set function G which is o-submodular and monotone, let the modular function mC(S)
[S] be defined as follows:
G((i, k)
G GGG @ ((1,))) CanG((i))) (((A))) g
Σ
OG
(8.R)E8
(K.4)888S
(8
(

---

Under review as a conference paper at ICLR 2023
If G' is a-submodular and monotone. it holds that G(S) I m [S] for all 8 S D (Durga et al. 2021). Note that when G is submodular, L... when a 1, the expression m gives us the existing modular upper bounds for submodular functions (Nemhauser et al. 1978 Iyer et al.
Algorithm 1 The PRESTO Algorithm
Require: Training data D, oc, K model architectures, Iterations.
1: Output: The learned parameters {0x k E [K]]. the partitioning of D: D UketaySL
2: S, + o...., 4
2013a Lyer & Bilmes 2012).
for E [K]
4: 02 <- INITPARAMS()
Outline of PRESTO (Alg. D. The goal of Al-
for all E do
gorithm 1 iteratively minimize m [S] with
6:
G({(i,k)]) + Train(fi);
respect in a majorization-miiiiizzn
7: end for
manner [Iyer et al., 2013a bj Durga et al. 2021]. Having computed that minimizes [8] in the iteration r 1, we set S = S for the iteration r and minimize mg m[(3) with re- spect to S.
8:
end for
9: for [Iterations] do for k: [K] do
10:
11:
3. - (i,k) i Sx
12:
0,. ( + Train (Se: Os)
We start with S = 0 (the empty partition) in line 2. Since we minimize mg G [3] with re-
13:
for E S, do
14:
G( k)) Train))rain(SS (Sx (z)
spect to S, the key components for calculation
15:
M Mi] + ag G G \ (()
are the last two terms of the RHS in Eq. ( i.e, G(f(i, ()))/ and G((i, k) S (ii 17:
16:
end for
for all i Sv. M [i][[] G(((,,
We pre-compute the former in lines 3-8 for all pairs (i,k). The complexity of this operation is (() We can compute 20:
18:
end for
19:
(i", k²) +- argmin, argmin, (M((()))
Se. S.. U2"
G((i,k) K) (V((i,k) k+) by finding the correspond- 21:
For all k: # k" =S 6 Sx (i+}
ing partitions { 3x3 and for each k. compute G(Sa) and G(S k)), (i, E S (lines 13-16). Finally, once we compute the modular upper bound mg (S). we minimize it
22:
+ (1, k) Se, k E [K]]
23:
end for
24: Return 0,,..., 0x.
subject to the matroid span constraint- thus, we get a partition. Since G is a modular function, we are guaranteed to find set which is in the base of the partition matroid. Finally, note that to evalu- ate G, we need to train the algorithm on the specific datapoints and partition, as encapsulated the Train() routine in Algorithm Here, Train(S, (S, 0x) trains the model for partition on the subset Sx for few iterations and, returns the estimated parameters Og and the objective G(Sk).
Approximation guarantee. The following theorem shows that and curvature KG. Algorithm an approximation guarantee (Proven in Appendix A.3))
a-submodular with a >> og.
Theorem 4. Given the set function G(S) defined in Eq. 44). let OPT be an optimal solution of the optimization problem (5) (or equivalently 77). If Algorithum [ returns & as the solution, then ( S G(OPT)/(11 - (G) [G) where Kg and og are computed using Theorem 2
EXPERIMENTS WITH REAL DATA
In this section, we provide a comprehensive evaluation of our method on four real world classi- fication datasets and show that our method is able to outperform several unsupervised partitioning methods and mixture models. Next, we evaluate PRESTO in the context of resource constrained deep learning methods, where we show that PRESTO is able to provide comparable accuracy with respect to state-of-ther thetthr-rrt models for network pruning and knowledge distillation. Appendix contains additional experimental results.
4.1 EXPERIMENTAL SETUP
Datasets. We experiment with four real world classification datasets, viz., CIFARIO (Krizhevsky 2009). PathMNIST (PMNIST) (Yang et al.] (202I). DermaMNIST (DMNiST) Yang et al.) 202
and SVHN (Netzer et al. 2011 Appendix D contains additional details about the datasets.
7

---

Under review as a conference paper at ICLR 2023
Accuracy P(ý g)
Macro Fl-score
Method
CIFARTO
PMINITT
DMNIST
SVHN
CIFARTO
PMNNST
DMINST
SVHN
Equal Kmeans
89,83
82.91
73.97
87.41
89.84
77.60
43.00
86.00
Kmeans++
89,50
82.64
73.72
87.02
89.48
77.20
44.30
85,50
Agglomerative
89.35
81.67
75.46
87.26
89.34
75.30
47.00
85,90
81.29
72.37
85.68
75.57
37.11
84,23
GMM
89.66
83.20
73.72
87.15
89.65
77.60
44.30
85,80
BGMM
73.72
87.35
75.82
44.30
85.90
Learn MLR
89.99
81,57
74,02
87.07
99.99
75.90
43.60
85.70
PRESTO
90.04
83.70
90.00
77.80
47.70
86.80
Table 1: Comparison of classification accuracy PP = y) and Macro-F1 score of PRESTO against three unsupervised partitioning methods Kmeans (Bennett et al. 2000), Kmeans+ (Arthur & Vassilvitskiii (006). Agglomerative (Müllner) (011)); Mixture of Experts(MoE)) (Shazeer et al, 2017) and three mixture models (GMM (Bishop & Nasrabadi] 2006). BGMM (Attias [999) and
Learn-MLR (Pal et al. 2022)) for all datasets. In all cases, we used exactly the same model ar- chitecture for each component, which is a ResNet18 network with reduced capacity (described in
Section 4.1 The best and second best results are highlighted in green and yellow respectively.
Implementation details. For all models, we extract features ø from fifth layer of pre-trained ResNet18 (He et al.) 2016). We design he, using neural architecture similar to ResNet18 with reduced capacity- specifically, starting with the sixth layer and until the eighth layer, we reduce the number of convolutional filters to 4 instead of 128, 8 instead of 256 and 16 instead of 512, respectively. In each case, we set the number of model components K using cross validation. We found K 4 for all datasets except DMNIST and K 3 for DMNIST. We use fully connected single layer neural network to model the additional classifier n that is used to decide the model component to be assigned further details in Appendix D
instance during inference. The experimental setup is discussed in
4.2 RESULTS
Comparison with unsupervised partitioning methods and mixture models. We first compare our method against several unsupervised partitioning methods and mixture models. In the case of unsupervised partitioning, we use clustering methods to slice the data into K partitions before the start of the training and then use these data partitions to train the K models. Specifically, we consider three unsupervised partitioning methods, viz., (1) Equal-Kmeans (Bennett et al., 2000) which is a K-means clustering method that ensures equal distribution of instances across different clusters, (2) Kmeans++ (Arthur & Vassilvitskii 2006). a variant of K-means method that uses a smarter initialisation of cluster centers. (3) Agglomeration clustering (Müllner) (011)) where the clusters are built hierarchically; (4) Mixture of (XoFrts(MoF) (Shazeer et al.) [017). a Sparsely- Gated Mixture-of- Experts layer (MoE), which selects sparse combination from a set of expert models using a trainable gating mechanism to process each input, and three mixture models, viz, (5) Gaussian mixture models (GMM) (Bishop & Nasrabadi 2006). (6) Bayesian Gaussian mixture model (BGMM) (Attias) 1999) and (7) Learn-MLR (Pal et al. 2022) which presents a version of the AM algorithm for learning mixture of linear regressions. We adapt this algorithm for classification as a baseline for PRESTO. Across all baselines, we employed exactly same set of model architectures. More implementation details of these methods is given in Appendix D
In Table 1 we summarize the results in terms of classification accuracy TP(1 = y) and Macro-F1 score. The Macro-FI score is the harmonic mean of the precision and recall. We make the following observations. (1) PRESTO demonstrates better predictive accuracy than all the baselines. (2) For three out of four datasets, the unsupervised partitioning methods outperform the mixture models. Note that there is no consistent winner among the partitioning methods. Thus, the best choice of the underlying partitioning algorithm changes across different dataset. However, since the optimal clusters obtained by our method are guided by supervised loss minimization, PRESTO performs well across all datasets.
Application on resource constrained learning. Given K, the number of partitions, PRESTO works with K lightweight models, which can be trained on different partitions of dataset, instead of training a very large model on the entire dataset. To evaluate the efficacy of PRESTO a resource constrained learning setup, we compare it against two model pruning methods, viz., (1) GraSP (Wang et al. 2020a), (2) SNIP (Lee et al. 2019) and two methods for knowledge distillation, viz., (3) KD (inton et al. 2015) (4) DGKD (Snn et al. 202T) in terms of both accuracy and GPU memory used during training. Pruning methods mask connections of larger model that relatively less
8

---

Under review as a conference paper at ICLR 2023
Method
Accuracy P(y y)
Training time GPU Usage (MiB)
CIFRRO
PMNIST
DMNIST
SVHN
CIFARIO
PMNIST
DMNIST
SVHI
GraSP
89.37
85.60
74.71
89.16
1107.21
1173.46
1173.16
1173.97
SNIP
88.24
85.28
73.77
89.36
863.73
1149.08
910.49
910.54
KD
90.06
75.32
72.37
89.08
272.00
250.69
250.56
251.81
DGKD
90.17
76.24
72.02
88,97
272.00
250.09
250.56
252.06
PRESTO
90.04
83.70
75.61
88.11
173.43
182.24
58.12
234.48
Table 2: Comparison of accuracy P(ù y) and maximum GPU memory used during training
(in mebibytes, MiB) between PRESTO and two model pruning methods, viz., GraSP (Wang et al. 2020a) and SNIP (Lee et al., 2019) and two knowledge distillation methods, viz., KD (Hinton al., 2015) and DGKD (Son et al. 202T on 20% held-out set. Numbers in green and yellow indicate the best and the second best method.
00.4
CIFARIO
90.2
Printo KO
Proto RD
00.0
Ponda
DOKD
80.0
DOKD
Accuracy
50.0
DOKD
60.8
KD
DUND
77.3
PAINIST
$9.8
CIPARIOO
77.5
PANNIT
160
190
220
250
280
73 170
195
220
245
270
10
15
20
26
30
78.0
12
15
GPU (in MiB)
GPU (in MIB)
GPU (iu MiB)
GPU (in MIB)
(a) Training
(b) Inference
Figure 3: Trade off between accuracy and maximum GPU memory for KD, DGKD and PRESTO during training (panel (a)) and inference (panel (b)) for CIFARIO and PMNIST.
important for training, whereas knowledge distillation methods start with high capacity model (teacher model) and then design a lightweight model (student) that can mimic the output of the large model with fewer parameters. For fair comparison, we maintain roughly similar (with 5% tolerance) number of parameters of each of the final lightweight models obtained by the pruning methods and of the student models given by the knowledge distillation methods; the total number parameters used in our approach is Σ dim(0x)
First, set K of PRESTO same in Table 1 This gives the number of parameters of PRESTO as 82690 for CIFAR1O, 54993 for PMNIST, 4T047 for DMNIST and 55138 for SVHN. In Table summarize the results for this setup. We observe that PRESTO consumes significantly lesser GPU memory than any other method across all datasets, whereas it yields the best accuracy for DMNIST. For DMNIST, it consumes an impressive 77% lesser GPU memory than the second most efficient method, viz. DGKD. The existing baselines use high capacity models, as reference models in pruning and teacher models in knowledge distillation, resulting in high training GPU memory consumption.
Next we vary the number of parameters p of PRESTO (through K) as well as the baselines E [62020, 124030] for CIFARIO, p E [41247, 82485] for PMNIST, p E [41047, 82087] for DM- NIST, p E [41356, 82702] for SVHN and probe the variation accuracy vs. maximum GPU memory used during both training and test. In Figure 3 we summarize the results for PRESTO and KD and DGKD- the two most resource efficient methods from Table 2 We make the following observa- tions. (1) PRESTO consumes significantly lower memory for diverse model size during both training and inference; and, (2) In most cases, as we change the model size, the accuracies obtained by the baselines vary widely whereas, model often shows only insignificant changes we change the model size.
CONCLUSION
We present PRESTO, novel framework of learning mixture models via data partitioning. Individual models specialising on a data partition, present good alternative to learning one complex model and help achieve better generalisation. We present joint form the partitions with good approximation guarantees. With our experiments we demonstrate that PRESTO achieves best performance across different datasets when compared with several mixture- models and unsupervised partitioning methods. We also present that PRESTO achieves best accuracy vs. memory utilisation trade-off when compared with knowledge distillation and pruning methods. Our work opens several areas for research, such as that of handling datasets with larger class imbal- ance, outlier detection, handling out of distribution data, etc.
optimization algorithm to

---

Under review as a conference paper at ICLR 2023
REPRODUCIBILITY STATEMENT
We provide zip file containing all the source code for PRESTO as well as the various baselines we use, as part of the supplementary material. The implementation details and machine configuration for our experiments are given in Section D of the Appendix. All our datasets are public and can be obtained easily. Users can download the datasets and use the provided code to reproduce the results presented in this paper.
REFERENCES
Dimitris Achlioptas and Frank McSherry. On spectral learning mixtures of distributions. In International Conference Computational Learning Theory, pp. 458-469. Springer, 2005.
Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameterization. In International Conference on Machine Learning. pp- 244-253. PMLR, 2018.
David Arthur and Sergei Vassilvitskii. k-means++: The advantages of careful seeding. Technical report, Stanford, 2006.
Hagai Attias. A variational baysian framework for graphical models. Advances in neural information processing systems, 12, 1999.
Ramakrishna Bairi, Rishabh Iyer, Ganesh Ramakrishnan, and Jeff Bilmes. Summarization of multi- document topic hierarchies using submodular mixtures. ACL, pp. 553-563, 2015.
Sivaraman Balakrishnan, Martin J Wainwright, and Bin Yu. Statistical guarantees for the em algo- rithm: From population sample-based analysis. The Annals of Statistics, 45(1):77-120, 2017.
Arindam Banerjee, Srujana Merugu, Inderjit S Dhillon, Joydeep Ghosh, and John Lafferty. Cluster- ing with bregman divergences. Journal of machine learning research, 6(10), 2005.
Mikhail Belkin and Kaushik Sinha. Polynomial learning of distribution families. In 2010 IEEE 51st Annual Symposium Foundations of Computer Science, pp. 103-112. IEEE, 2010a.
Mikhail Belkin and Kaushik Sinha. Polynomial learning of distribution families. In 2010 IEEE 5 Annual Symposium on Foundations of Computer Science, pp. 103-112. IEEE, 2010b.
K.P. Bennett, P.S. Bradley, and Demiriz. Constrained k-means clustering. Technical Report NSS--RR-0000-65, May 2000.
Christopher M Bishop and Nasser M Nasrabadi. Pattern recognition and machine learning, vol- ume 4. Springer, 2006.
Christos Boutsidis, Petros Drineas, and Malik Ismail. Near-optimal coresets for least- squares regression. IEEE transactions on information theory. 59(10))6880-689) 2013.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neclakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Sébastien Bubeck and Mark Sellke. A universal law of robustness via isoperimetry. Advances in Neural Information Processing Systems, 34, 2021.
Trevor Campbell and Tamara Broderick. Bayesian coreset construction via greedy iterative geodesic ascent. In International Conference on Machine Learning, pp. 698-706, 2018.
Ronan Collobert, Samy Bengio, and Yoshua Bengio. A parallel mixture of syms for very large scale problems. Advances in Neural Information Processing Systems, 14, 2001.
Yehuda Dar, Vidya Muthukumar, and Richard G Baraniuk. A farewell to the bias-varianc tradeoff? an overview of the theory of overparameterizee machine learning. arXiv preprint 10Xiv:2109.02355. 2021.
10

---

Under review as a conference paper at ICLR 2023
Sanjoy Dasgupta. Learning mixtures of gaussians. In 40th Annual Symposium Foundations of Computer Science (Cat. No. 99CB37039), Pp. 644. IEEE, 1999.
Arthur P Dempster, Nan M Laird, and Donald B Rubin. Maximum likelihood from incomplete data via the em algorithm. Journal of the Royal Statistical Society: Series (Methodologicul), 39(1): 1-22, 1977.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hier- archical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pp. 248 248-255, 2009. doi: 10. 10.11009.2..2 5206848.
Jacob Devlin, Ming- Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.048005, 2018.
José G Dias, Jeroen K Vermunt, and Sofia Ramos. Mixture hidden markov models in finance re- search. In Advances in data analysis, data handling and business intelligence, pp. 451-459. Springer, 2009.
Sivasubhramanian Durga, Rishabh Iyer, Ganesh Ramakrishnan, and Abir De. Training data subset selection for regression with controlled generalization error. Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Rexearch, pp. 9202-9212. PMLR, 18-24 Jul 2021.
Jack Edmonds. Submodular functions, matroids, and certain polyhedra. In Combinatorial Opti- mization-Ein Eureka, You Shrink!, pp. 11-26. Springer. 2003.
Marwa El Halabi and Stefanie Jegelka. Optimal approximation for unconstrained non-submodular minimization. In International Conference on Machine Learning, pp. 3961-3972. PMLR, 2020.
Susana Faria and Gilda Soromenho. Fitting mixtures of linear regressions. Journal of Statistical Computation and Simulation, 80(2):201-225, 2010.
Zhouyu Fu and Antonio Robles-Kelly. On mixtures of linear svms for nonlinear classification. In Joint IAPR International Workshops on Statistical Techniques in Pattern Recognition (SPR) and Structural and Syntactic Pattern Recognition (SSPR). pp. 489 489-499, 499. Springer, 2008
Khashayar Gatmiry and Manuel GodezzRodriguezz. Non-submodular function maximization subject to matroid constraint, with applications. arXiv preprint 18 2018.
Amir Gholami, Kiseok Kwon, Bichen Wu, Zizheng Tai, Xiangyu Yue, Peter Jin, Sicheng Zhao, and Kurt Keutzer. Squeezenext: Hardware-aware neural network design. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR) Workshops, June 2018
Abolfazl Hashemi, Mahsa Ghasemi, Haris Vikalo, and Ufuk Topcu. Submodular observation selec- tion and information gathering for quadratic models. arXiv preprint arXiv: wtt.:/995599999 2019.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp- 770-778, 2016.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network, 2015.
Samuel B Hopkins and Jerry Li. Mixture models, robustness, and sum of squares proofs. In Pro- ceedings of the 50th Annual ACM SIGACT Symposium on Theory of Computing, pp. 1021-1034, 2018.
Andrew G. Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam. Mobilenets: Effficient convolutional neural networks for mobile vision applications, 2017. URL httttttt///xxiii. org. abs/ 1704. 04861
Forrest N. landola. Song Han, Matthew W. Moskewicz, Khalid Ashraf, William J. Dally, and Kurt Keutzer. Squeezenet: Alexnet-level accuracy with 50x fewer parameters and 0.5mb model size, 2016. htl.sttts: www.:oorg/bbboo abba//////// 673607360

---

Under review as a conference paper at ICLR 2023
Rishabh Iyer and Jeff Bilmes. Algorithms for approximate minimization of the difference between submodular functions, with applications. arXiv preprint arXiv:1207.0560, 2012.
Rishabh Iyer and Jeff Bilmes. Polyhedral aspects of submodularity, convexity and concavity, arXiv preprint SEIG::1506.072222 2015.
Rishabh Iyer, Stefanie Jegelka, and Jeff Bilmes. Fast semidifferential-aseed submodular function optimization: Extended version. ICMI., 2013a.
Rishabh Iyer, Stefanie Jegelka, and Jeff Bilmes. Monotone closure of relaxed constraints in sub- modular optimization: Connections between minimization and maximization: Extended version. In UAL. Citeseer, 2014.
Rishabh K Iyer, Stefanie Jegelka, and Jeff A Bilmes. Curvature and optimal algorithms for learning and minimizing submodular functions, In NeurIPS, 2013b.
Yuang Jiang, Shiqiang Wang, Victor Valls, Bong Jun Ko, Wei-Han Lee, Kin K Leung, and Leandros Tassiulas. Model pruning enables efficient federated learning edge devices. arXiv preprini 0rXiv:1909.12320 2019.
Michael Jordan and Robert A Jacobs. Hierarchical mixtures of experts and the em algorithm. Neural computation, 6(2):181-214, 1994.
Esmail Josh Levy-Kramer, Matt Klaber. k-means-consttrined. https://github: com/
joshik/k-means constrainee 2022.
Adam Tauman Kalai, Ankur Moitra, and Gregory Valiant. Efficiently learning mixtures of two gaussians. In Proceedings of the forty-second ACM symposium Theory of computing, pp- 553-562, 2010.
Vishal Kaushal, Rishabh Iyer, Suraj Kothawade, Rohan Mahadev, Khoshrav Doctor, and Ganesh Ramakrishnan. Learning from less data: A unified data subset selection and active learning framework for computer vision. In 2019 IEEE Winter Conference on Applications of Computer Vision ( (WACV), pp. 1289-1299. IEEE, 2019.
Krishnateja Killamsetty, Durga Sivasubramanian, Baharan Mirzasoleiman, Ganesh Ramakrishnan, Abir De, and Rishabh Iyer. Grad-match: A gradient matching based data subset selection for efficient learning. arXiv preprini arXiv:2 2021a.
Krishnateja Killamsetty, Durga Subramanian, Ganesh Ramakrishnan, and Rishabh Iyer. Glister: A generalization based data selection framework for efficient and robust learning. In AAAI, 2021b.
Katrin Kirchhoff and Jeff Bilmes. Submodularity for data selection in machine translation. In Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 141- 2014.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
Jeongyeol Kwon, Wei Qian, Constantine Caramanis, Yudong Chen, and Damek Davis. Global convergence of the em algorithm for mixtures of two component linear regression. In Conference on Learning Theory, pp. 2055-2110. PMLR, 2019.
Namhoon Lee, Thalaiyasingam Ajanthan, and Philip Tor. SNIP: SINGLE-SHER SHOT NETWORK PRUNING BASED ON CONNECTION SENSITIVITY. In International Conference on Learn- ing Representations, 2019. URL /tpptttte////peeeeiieee hwwp//eeeprreeeeeeeeeeeeeeeeee.e net. forum wtwp:ummiiiiiiiiaae AccX
Bailin Li, Bowen Wu, Jiang Su, and Guangrun Wang. Eaglecye: Fast sub-net evaluation for efficient neural network pruning. In European conference on computer vision, pp. 639-654. Springer, 2020.
Yuanzhi Li and Yingyu Liang. Learning mixtures of linear regressions with nearly optimal com- plexity. In Sébastien Bubeck, Vianney Perchet, and Philippe Rigollet (eds.), Proceedings of the 31 Conference On Learning Theory, volume 75 of Proceedings of Machine Learning Research, pp. 1125-1144, PMLR, 06-09 Jul 2018. h https: http:::://ppoceedinng.. mir. www. a.l.ppess/v75/
111 118b. html
12

---

Under review as a conference paper at ICLR 2023
Roman Liesenfeld. A generalized bivariate mixture model for stock price volatility and trading volume. Journal of econometrics, 104(1):141-178, 2001.
Lin, Yongming Rao, Jiwen Lu, and Jie Zhou. Runtime neural pruning. Advances neural information processing systems, 30, 2017.
Tao Lin, Sebastian U Stich, Luis Barba, Daniil Dmitriev, and Martin Jaggi. Dynamic model pruning with feedback. arXiv preprint arXiv:2006.07253, 2020.
Chaoyue Liu, Libin Zhu, and Mikhail Belkin. Toward theory of optimization for over- parameterizze systems of non-linear equations: the lessons of deep learning. arXiv preprint arXiv:2003.00307. 2020.
Yinhan Liu, Myle Out, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1692, 2019.
Yuzong Liu, Rishabh Iyer, Katrin Kirchhoff, and Jeff Bilmes. Svitchboard and fisver i: High- quality limited-complxxity corpora of conversational english speech. In Sixteenth Annual Confer- ence of the International Speech Communication Association, 2015.
Stuart Lloyd. Least squares quantization in pem. IEEE transactions on information theory, 28(2): 129-137, 1982.
Mario Lucic, Matthew Faulkner, Andreas Krause, and Dan Feldman. Training gaussian mixture models at scale via coresets. The Journal of Machine Learning Research, 18(1):5885-5900 2017.
Ningning Ma, Xiangyu Zhang, Hai-Tao Zheng, and Jian Sun. Shufflenet v2: Practical guidelines for efficient cnn architecture design. In Proceedings of the European conference on computer vision (ECCV), pp. 116-131, 2018.
J MacQueen. Classification and analysis of multivariate observations. In 5th Berkeley Symp. Math. Statist. Probability, pp. 281-297, 1967.
Baharan Mirzasoleiman, Jeff Bilmes, and Jure Leskovec. Coresets for data-efficient training of machine learning models, 2020a.
Baharan Mirzasoleiman, Kaidi Cao, and Jure Leskovec. Coresets for robust training of neural net- works against noisy labels. arXiv preprint arXiv:2011.07451, 2020b.
Daniel Müllner. Modern hierarchical, agglomerative clustering algorithms. CoRR, abs// 23782
2011. URL http: www.uuniittt.... http:://ww................/////r wttp:/aaaaaaaoooeororeorro htt.11199...m.m..mmmmmmm
abs-1109-2378
George L Nemhauser, Laurence A Wolsey, and Marshall L Fisher. An analysis of approximations for maximizing submodular set functions-4, Mathematical programming, (6()):2652294, 1978.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Reading digits in natural images with unsupervised feature learning. 2011.
Kelley Pace and Ronald Barry. Sparse spatial autoregressions. Statistics Probability Letters, 33 (3):991-297, 1997.
Soumyabrata Pal, Arya Mazumdar, Rajat Sen, and Avishek Ghosh. On learning mixture of linear regressions in the non-realizable setting. In International Conference on Machine Learning, pp. 17202-17220. PMLR, 2022.
Wei Pan, Jizhen Lin, and Chap Le. A mixture model approach to detecting differentially expressed genes with microarray data. Functional & integrative genomics, 313):117-124, 2003.
Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Pretten- hofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Coumapeau, M. Brucher, M. Perrot. and E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research, 12:2825-2833 2011.
13

---

Under review as a conference paper at ICLR 2023
David Rau. pharsely-gated mixture-of-experts pytorch implementation, 2019.
Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chied Chen. Mo- bilenetv2: Inverted residuals and linear bottlenecks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 4510-4520, 2018.
Arora Sanjeev and Ravi Kannan. Learning mixtures of arbitrary gaussians. In Proceedings of the thirty-third annual ACM symposium on Theory of computing, pp. 247-257725 2001.
Alexander Schrijver al. Combinatorial optimization: polyhedra and efficiency, volume 24. Springer, 2003.
Noam Shazeer, *Azalia Mirhoseini, *Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hin- ton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of- experts layer. In International Conference on Learning Representations, 2017. URL https: //openzeview. nwt http//orrroiiidclcccdddaa
Wonchul Son, Jaemin Na, Junyong Choi, and Wonjun Hwang. Densely guided knowledge distilla- tion using multiple teacher assistants, 2021.
Nicolas Städler, Peter Bühlmann, and Sara Van De Geer. 11 I1-penalization for mixture regression models. Test, 19(2):009 256,, 2010.
Jingtong Su, Yihang Chen, Tianle Cai, Tianhao Wu, Ruiqi Gao, Liwei Wang, and Jason D Lee. Sanity-checking pruning methods: Random tickets can win the jackpot. arXiv preprint arXiv:2009.11 1094, 2020.
Mingxing Tan and Quoe Le. Efficientnet: Rethinking model scaling for convolutional neural net- works. In International conference on machine learning, pp. 6105-6114. PMLR, 2019.
Chaoqi Wang. Guodong Zhang, and Roger Grosse. Picking winning tickets before training by
preserving gradient flow. In International Conference on Learning Representations, 2020a. URL
httts:://openreve iew. net. wwt.://wwweeteeeteeeeeeeeeeeee f
Yulong Wang. Xiaolu Zhang, Lingxi Xie, Jun Zhou, Hang Su, Bo Zhang, and Xiaolin Hu. Pruning from scratch. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 12273-12280, 2020b.
Kai Wei, Rishabh Iyer, and Jeff Bilmes. Fast multi-stage submodular maximization. In International conference on machine learning, pp. 1494-1502. PMIR, 2014a
Kai Wei, Yuzong Liu, Katrin Kirchhoff, and Jeff Bilmes. Unsupervised submodular subset selec- tion for speech data. In 2014 IEEE International Conference Acoustics, Speech and Signal Processing (ICASSP), pp. 4107-4111. IEEE, 2014b.
Lei Xu and Michael Jordan. On convergence properties of the em algorithm for gaussian mixtures. Neural computation, 8(1):129-151) 1996.
Jiancheng Yang, Rui Shi, Donglai Wei, Zequan Liu, Lin Zhao, Bilian Ke, Hanspeter Pfister, and Bingbing Ni. Medmnist v2: A large-scale lightweight benchmark for 2d and 3d biomedical image classification. arXiv preprint arXiv:2008, 2021.
Xinyang Yi, Constantine Caramanis, and Sujay Sanghavi. Alternating minimization for mixed linear regression. In International Conference Machine Learning, pp. 613-621. PMLRR, 2014.
Haifeng Zhang and Yevgeniy Vorobeychik. Submodular optimization with routing constraints. In Proceedings of the AAAI conference on artificial intelligence, volume 30, 2016.
Xiangyu Zhang, Xinyu Zhou, Mengxiao Lin, and Jian Sun. Shufflenet: An extremely efficient
convolutional neural network for mobile devices. In 2018 IEEE/CVF Conference on Computer Vision and Paterrn Recognition, pp. 1848-6854 2018. doi:
14

---

Under review as a conference paper at ICLR 2023
Appendix
A PROOFS OF THE TECHNICAL RESULTS IN SECTION 3
A.1 FORMAL DISCUSSION ON MATROID
Definition 5. A matroid is combinatorial structure M := (V,I) defined on a ground set V and family of independent sets I 2 which satisfies two conditions.
(1) It'S C T and T I, then $ E I. (2) It'S E and T E I and [T] > |S], then there exists a e E T so that SU {e) EI.
From the above definition, is clear the all the maximal independent sets have same cardinality. A maximum independent set is called the base of the matroid.
A.2 MONOTONICITY AND C-SUBMODULARITY OF G
Here we prove the claims Theorem 2 We repeat the theorem for convenience
Theorem 2 Given the set function G(S) defined in Eq. (4) and the individual regularized loss
functions E introduced in Eq. D. we define: Emin maXk,; Eigen V E(ho. (ax), 31), Luin
min ming, [E(ho, (zx,), yt) + x1(0x117] and Emin- = m ming, [E(ho, (2x)), yc) + \110x11?)
Then, we have the following results:
is monotone non-decreasing S. (2) o-submodularity: If E(ho, (z).v) is L-Lipschitz, for all k E (1,.., K} and that the regularizing coefficient À satisfies: X > -Emin function G(S) is susubmodular with
(1) Monotonicity: The function G(S) defined in Eq.
miin
Emin
21
AL2
X+50.5cmmn
(1+00.5mm
(3) Generalized curvature: The generalized curvature Ka (S) for any set S is given by: (G(()) KiG mmin/ E
Proof. (1) (Proof of monotonicity) We define the regularized g(O,, S) = Eies (ho (wi), yi) + x| First, we show that (S (()) {)) 0. To that aim, note that, g( (SSUU)) SU (D - (S))
(S (j)). SU (Л)) (SU ((0) -())) (10)
(SO)
Now, since (S) argmine, S), we have (SU())))) 2 ((0)
From Eq. 10
this leads the following inequality:
( U (i)), (i)) (0) 2 ( ()), (5)) 2 Amin >0
(11)
Now, if we include new element (j,t) into S, then will only change the loss g(0 (S), among all model components. Thus we have,
G(S U ((.0)) - G(S) 2 (0) (SU (J), (SU (i)) ((0
(2) (Proff of a-submodularity) First, we try to show a-submodularity of (S), S). Hence, we first bound the following ratio:
g US -
(12)
g ( (f)). (TO)
Eq. directly provides bound on the numerator of this ratio:
Next, we try to upper bound (() -(( ((O; Luin
g (
(13)
We note that:
g (T -()
(14)
Now, since ( ( (i)) argming, 50x, TU (11). we have:
( ( -
(15)
15

---

Under review as a conference paper at ICLR 2023
Next, note that 1111://777711113---000000000-110 which gives us:
-0 (i))
(16)
This leads us to bound the ratio in Eq. (2) as follows:
( -
(17)
(T -
(())
g(0;(T);
U (i)) ())
(18)
((0 + (T), (i}) - ((0
To bound the denominator, we note that:
(T). (T) (S (5)), (1) x (z)) - ((hex(suu)(())).)) + (T))
Now, we note that
T
2 g(O((T), - 12 + (oo((ho((z)).n)) O
o=0
1
2 (T)
Emin + Emin
1110 - 1| < 0
(19)
Thus, we have: ||0x|| S x+ Emgin
Putting this in Eq.
we have:
((0((T), (i}) ((OT(S (b) < L (T)(T) ( (U) A (T)]
(20)
21
AL
+
(21)
X + Emin
(A+
Emaa )
Thus, replacing the above quantity in Eq.
18
we have:
lin
(22)
Cuin
0. )
(3) (Proof of the bound on Generalized curvature) From the definition of curvature, (0) (5 (a)), {a})
(S) min
(23)
.EV g(O(((a)), {a})
From Eq. 11
((0 ( S\ {a}) 2 Lmin
(24)
g (O (a)) < max |^| + I(Se, (Ia), No)] Emin
(25)
Thus,
(g) S 1 Luin
(26)
Lmin
we add an element a = (j. t) to S only ((,, will be changed among the component models. Thus,
G(S) - G(S) (at)
(S) = 1 min
(27)
aeV
G(a)
g(0r ((0 (a)), (a))
min
aeV
((r((a))
mmin
Luin
16

---

Under review as a conference paper at ICLR 2023
A.3 APPROXIMATION GUARANTEES
We next prove the approximation bound for Theorem
Theorem 4). If the function G is uG-submodular and has a curvature KG. Algorithm obbtins an
approximation guarantee of aed[+((D)-1)0--aaaar S og(1-Nc) assuming there exists perfect training oracle in Lines (6, 14)
[2]
Proof. From the definition of a-submodularity, note that agG(S) S Eies G(i). Next, we can obtain the following inequality for any & using weak submodularity:
G( G(k) 2 ag (G((S))
(28)
JEBYA
We add this up for all k E S and obtain:
[S) Σ
G(k) 2 aG Σ Σ (G((S))
kEB
KEB JEBY&
ac((8|-1 Σ G(K(S\A)
(29)
KgB
Finally, from the definition of curvature, note that G(K[S\k) I (1 Ky G(k). Combining all this together, we obtain:
S G N (1+ ( 1)) Σ G())
(30)
jES
which implies:
S
ΣG() <
G(S)
(31)
1+ ac xcx) 1)
JES
Combining this with the fact that ag G(( $ E iiS G(i),
obtain that:
1
S
G(S)
G(i)
ag (1 + (cc 1)) G(S)
(32)
aG
Kg)
iE3
Note that [SI((1 + aG KG) ( 1) S 1/ - kG) so we just use this factor in the ap- proximation bound. The approximation guarantee then follows from some simple observations. In particular, given an approximation
m G (S)
Σ G(I SES
(33)
OG
which satisfies G(S) < (S) BcG(S), claim that optimizing ma essentially gives Bo approximation factor. To prove this, let S be the optimal subset, and $ be the subset obtained after optimizing m The following chain of inequalities holds:
G < m ($) S m ((') < GS"
(34)
This shows that S is 3o approximation of St. Finally, note that this is just the first iteration of PRESTO, and with subsequent iterations, PRESTO is guaranteed to reduce the objective value since we only proceed there reduction in objective value.
17

---

Under review as a conference paper at ICLR 2023
B COMPARISON WITH EM ALGORITHM
General EM algorithm usually considers a prior and computes the membership probability of each datapoint at each iteration. It makes assumption about the generative mechanism of the data well as the prior about the cluster probabilities. Such probabilistic approach naturally leads one to develop EM algorithm.
On the other hand, our method makes no assumption about the data as well as cluster membership. We do not make any probabilistic assumption which naturally led to develop set optimization problem- we do not make use of any "continuous" or "soft" scores on cluster membership. Thus, method is MM algorithm which is iterative set optimization algorithm, which is function- ally very different from EM algorithm.
C DIFFERENCE WITH MIXTURE OF LINEAR REGRESSION
All our experiments consider classification tasks. Therefore, we modify Learn-MIR (Pal et al. 2022) mixture of classifiers. However, note that the optimization problem formulation of the mixture of linear regression (or classification) can be viewed as special case of the problem formulation of our method.
In most cases, the existing methods make assumptions about the generative mechanism of the cluster membership (e.g., some definite prior) and resort to an algorithm which makes an involved use of that assumption and is significantly tailored to particular task (mixture of regression or classifica- tion).
In contrast, our framework operates on generic framework and does not make any specific assump- tion about the data or the cluster membership. Moreover, the algorithm very different from what is used by the usual mixture model learning algorithms.
D ADDITIONAL EXPERIMENTAL SETUP
D.1 DATASET DETAILS
We experiment with four real world classification datasets in experiments,
CIFARIO (Krizhevsky] 2009) has images of 10 different real-world objects.
SVHN (Netzer et al.) 2011) has real-world images for the 10 digits classification task
PathMNIST (PMNIST) Yang et al.) 2021) has images of colorectal cancer histology
slides.
DermaMNIST (DMNIST) (Yang et al. 2021 has dermatoscopic images of common pig-
mented skin lesions.
Table 4 lists the further details of these datasets. We see that CIFAR10 is a balanced dataset where PathMNIST, DermaMNIST and SVHN have large skew in their label proportions.
Dataset
#Classes
#Instances
#Train
#Test
#Avg points per class
CIFAR10 10
10
60,000
50,000
10,000
5,000 + 0
SVHN
10
99,289
73,257
26,032
7325.7 + 2800.66
PMNIST
9
107,180
89,996
7,180
9,999.56 + 1592.4
DMNIST
7
10,015
7,007
2,005
1,001 + 1,530.9
Table 4: Dataset statistics
Using a ResNet18 (He et al. 2016) model whose weights are initialised with weights obtained while training with Imagenet (Deng et al. 2009) dataset, get feature embedding of these datasets. For CIFAR1O dataset we extract embeddings from the fourth last layer of pretrained ResNet18 and for other datasets extract embeddings from the fifth last layer.
18

---

Under review as a conference paper at ICLR 2023
D.2 MODEL ARCHITECTURE USED FOR OUR METHOD, THE UNSUPERVISED PARTITIONING METHODS AND THE MIXTURE MODELS
In case of our method, the unsupervised partitioning methods and the mixture models (all methods in Table | and PRESTO in (2. we use a relatively light model. This model is exactly same in terms of architecture design to the final five layers of ResNet18 for PMNIST, DMNIST and SVHN, final four layers of ResNetl 8 for CIFAR10. This model design is chosen because we use extracted embeddings for training. The model reduces the convolution filters various junetures. The fifth last layer has filter of size 128 which is replaced with filter of size 4. Similarly, the filter at fourth last layer is changed from 256 to 8, filter third last layer is changed from 512 to 16. We train an additional classifier. This classifier is single layer fully connected feed forward network that aids in predicting the final class for each instance.
D.3 IMPLEMENTATION DETAILS FOR PRESTO
Number of partitions K. Number of partitions (K) is found for various datasets by performing an analysis of classification accuracy v/s K in the validation. Table shows K for each dataset.
Initialization. PRESTO initialised using one of the clustering techniques depending on the dataset. |
We also try to ensure that the clusters are of approximately the same size. We impose [K] constraints for each of the bin with some leeway. The leeway granted for every dataset is also tunable hyperparameter.
size
Final selection of (i, k) from the matrix M. The values of the matrix M computed in line 19 of Algorithm [ are supplemented with euclidean distance (D) of point the existing bin center. The final criterion for binning becomes M eD. Here, e = 3e - 3 for CIFARIO, SVHN and PMNIST. However, for DMNIST, we have 6e 3
Optimizer. We use an Adam optimizer with the Cross Cross-Entropy Loss to train the PRESTO models and the classifier. The learning rate set to 0.01. Table JabbleSjencapsulattee the various hyperparameters discussed for the datasets. We set the value of the regularizer À le
Dataset
K
Initialization
Leeway
CIFARTO
4
Fqual-KMeans
0
SVHN
4
Equal- Means
3000
PMNIST
4
Bqual-KMeane
3000
DMNIST
3
KMeans++ 500
Table 5:
Dataset specific hyperparameters for PRESTO
In Algorithm 1 Train(-) train the models on the respective data partitions for 10 epochs. We set Iterati ons 30. The additional classifier n is also trained for 10 epochs.
Table I mentions various baselines. Details about each baseline is covered in the following points.
D.4 IMPLEMENTATION DETAILS ABOUT THE UNSCPERVISED PARTITIONING METHODS AND MIXTURE MODELS
Equal-Kmeans.
We use the (Josh Levy-Kramer 2022) implementation of Constrained K means clustering. We run the algorithm for K clusters, and constrain the size of a cluster to be within IN K way parameter. A Kmeans++ initialization is used for selecting the initial centroids. We keep the number of times the algorithm will be run with different centroid seeds as n_init 10. The final result is the best output of n_init consecutive runs in terms of inertia. The algorithm runs either max of 300 iterations or convergence, whichever occurs earlier. The convergence criterion is defined as the Frobenius norm of the difference in the cluster centers of two consecutive iterations becoming less than 1e 4.
Leeway,
|N| | + Leeway i.e. roughly equally sized clusters with tolerance equal the Lee-
Kmeans++. We the scikit-leam (Pedregosa al. 2011 implementation of Kmeans++ with the following parameters: The number of clusters is set K. A Kmeans++ initialization is used for selecting the initial
19

---

Under review as a conference paper at ICLR 2023
centroids. We keep the number of times the algorithm will be with different centroid seeds as n_init = 10. The final result is the best output of n_init consecutive runs in terms of inertia. The algorithm runs either max of 300 iterations convergence, whichever occurs earlier. The convergence criterion is defined as the Frobenius norm of the difference in the cluster centers of two consecutive iterations becoming less than 1e 4.
bottom-up hierarchical clustering approach. We use the scikit-lear (Pe- dregosa et al. 2011) implementation of Agglomerative We use the Euclidean distance metrie to calculate distance between data points. Initially each data point is in its own cluster. The clusters are subsequently merged based on linkage criterion. We use the ward linkage criterion, which minimizes the sum of square differences within each cluster. The number of clusters is set as K.
Agglomerative.
GMM. We use the scikit-learn (Pedregosa et al. 2011 implementation of GMM with the number of components = K. The GMM weights are initalized according to KMeans initialization. The convergence threshold for the EM iterations is set as le -3
BGMM. We use the scikit-learn (Pedregosa et al. 2011) implementation of BMM with the number of components = K. The BGM weights are initialized according to KMeans initialization. The convergence threshold for the EM iterations is set as le - 3. The weight concentration prior is set as dirichlet process.
LeaRn-MLR. We adapt the original paper (Pal et al. 2022) for the purpose of classification. The initial partitioning is random and the we do not have a leeway when the data is re-partitioned. The e parameter which gives contribution to the euclidean distance (D) is set to 0. An Adam optimizer is used with Cross-Entrope loss. The learning rate is set to 0.01 and the regularizer X 1e 4.
Mixture of Experts. We the Sparse layer described in (Saazeer et al. 2017 for classi- fication. The expert networks have the architecture described in Section D.2] We adapt the (Rau 2019 implementation for this purpose. We use Noisy Top-K Gating as described in (Shazeer al., 2017 which balances the number of training examples each expert receives. We set the number of experts used for each batch element, k =44
D.5 IMPLEMENTATION DETAILS FOR CONBURRCECCONTAANEDDD EXPERIMENTS
We evaluate PRESTO in resource- constreecn learning setting by comparing with model pruning and knowledge distillation methods. The specifications for them are mentioned in the following points.
SNIP Lee et al. 2019). SNIP is a pruning at initialization method, that given a large reference desired sparsity level (i.e. number of non-zero weights). Pruning is done prior to training, in data-dependent way based on the loss function at a variance scaling initialization, to create a sparse subnetwork which then used for inference. We adapt the implementation from (Su et al. 2020) to experiments using SNIP. We train a large reference model, with the same architecture as the PRESTO models described in Section D.2] The reference model has the convolution filter size the fifth last layer 256, the fourth last layer as 512 and at the third last layer as 1024. The pruning ratio is set so that the number of nonzero parameters in the reference network after pruning are equal to that of the corresponding PRESTO model in all experiments. The models are trained on using an Adam optimizer and the Cross-Entropy loss with learning rate of 0. The reference models are trained for 300 epochs, and the subsequent pruned models are trained for total of 20 epochs till convergence.
network prunes network connections
GraSP (Wang al., 2020a). GraSP is pruning at initialization method similar to SNIP and also learns smaller subnetwork given reference network, which can subsequently be trained independently. The gradient norm after pruning the pruning criterion in GraSP, and those weights whose removal will result in least decrease in the gradient norm after pruning are pruned. We adapt the implementation from (Su et al. 2020) to run experiments using SNIP. We train a large reference model, with the same architecture as the PRESTO models described in Section D.2 The reference model has the convolution filter size at the fifth last layer as 256, at the fourth last layer as 512 and at the third last layer as 1024. The pruning ratio in all experiments is set that the number of nonzero parameters in the reference network after pruning are equal to that of the corresponding PRESTO model in all experiments. The models trained using Adam optimizer and the
20

---

Under review as a conference paper at ICLR 2023
Cross-Entropy loss with learning rate of 0.1. The reference models are trained for 300 epochs, and the subsequent pruned models are trained for a total of 20 epochs till convergence.
KD (Hinton et al. 2015). We train two models, teacher model and student model. The architec- ture of these models same as the PRESTO models described in Section [.2] The teacher model is heavier in size, the convolution filter size at the fifth last layer is 128, at the fourth last layer is 256 and at the third last layer is 512. On the other hand, the lighter student model has convolution filter size set 12,16 and 32. The models are trained on using Adam optimizer and the Cross-Entropy loss. The KD Temperature hyperparameter is set to 5. The teacher model is trained and the loss of the student model modified to include contribution from the teacher model. The multiplicative hyperparameter controlling the contribution of the teacher model is set to 0.9.
DGKD (Son al. 2021). This method requires to train models, a teacher model, a TA model and a student model. The teacher model is the heaviest, having the convolution filters set to 128, 256 and 512. The TA model is slightly lighter with filter size of 32, 64 and 128. The student model the lightest with filter size of 12, 16 and 32. The temperature hyperparameter is 5 and the factor controlling the contribution of the teacher and TA model is set at 0.9. We train the teacher and TA models independently. As in KD, the loss function of the student model is tweaked to consider a contribution from the teacher and the TA models.
Machine configuration. We performed our experiments computer system with Ubuntu 16.04.6 LTS, an i-7 with 8 cores CPU and a total RAM of 528 GBs. The system had a single Titan RTX GPU which was employed in our experiments.
21

---

Under review as a conference paper at ICLR 2023
80
R
11
07
+1
07
80
R
11
80
W 11
BV
09
09
1
09
09
09
09
09
09
(a) True Partitions
(b) Epoch r = 10
(c) Epoch r 20
(d) Epoch r = 50
Figure Snapshot of the true partition (panel (a)); and, the snapshots of the trained models {he 1 [K]] and the partitions {.S& k [K]} predicted by PRESTO (panels (b)-(d) on the syn- thetic dataset during progression of Algorithm The dataset is generated using a degenerate mixture model with K 4 components. Here, an instance (æ,, y:) belongs to exactly one of the four sets {Sk l E [4]] with probability 1, where St are defined in Eqs. (35 (38). Each model component is an SVM, i.e., he, (a) = w x + by and E(g. y) = (1 by)+. We observe that as r increases, PRESTO becomes more and more accurate in assigning instance the correct mixture component and finally, at r = 50, it is almost able to recover the true mixture component- the ground truth assignments of the instances (panel (a)) is extremely close the final assignments (panel (d)).
E EXPERIMENTS ON SYNTHETIC DATA
In this section, we experiment with synthetically generated instances from a latent degenerate mix- ture distribution and show that, Algorithm can accurately recover the partitions which correspond to the true mixture components.
E. EXPERIMENTAL SETUP
Dataset generation. We generate D = 20000 examples ((x),))) vee where X' [-1.73, 1.73) x [-1.17, 2.59] and y = {-1, +1}- The instances (æ,, y) belong to of the K 4 sets, viz,
S7, ST. defined as follows:
Si = ((.,y) [x[11] x/[] 2y 27 0 [- 1.73, 0.86 x (61.16, 0.04]]
(35)
S2 {(x. v) 20(x(1]3) - 200z[1] + 414 + x1 2y A x E [- 0.87 0.01] x [-1.15, 2.59]
(36)
5g = ((x, v)
1r + +
2y E 0.01, 0.87] x [-1.17, 2.56)
(37)
S7 ((x) v)
[x[1]] x=2] - 2y 43 = 0 [0886, 1.73] x [-1.17, 0.04]]
(38)
Panel (a) of Figure 6 shows scatter plot of (z,y). The instances are homogeneously distributed
across all components, i.e., | 5000.
Choice of fo, and L. We consider K = 4 linear support vector machines hx) w x: + by with O₂ = ( b% set l as the margin based hinge loss, i.e., f(ho, y) = (1-y(wua., +b )) )- Then we apply PRESTO (Algorithm D to simultaneously learn the parameters {(w,, bx) < k S 4) and the partitions S₁ such that D = S In addition to the classification accuracy, we also measure the aggregated error in predicting the correct mixture component which is defined as Errs = E.g 1 E S7. E St,j E Sithi t * t'). Thus, 0 computes the number of pairs that belong to the same latent mixture component Sj, but are predicted to have different mixture components by PRESTO.
E.2 RESCLTS
In Figure 6 we plot different snapshots of the models he. and the partitions D = D= Uf St. as PRESTO progresses during different iterations r (line 9 in Algorithm Dp. We make the following observations: (1) As r increases, the classification accuracy increases as well as mixture component prediction error Erry decreases. For r = 50, we observe that Errs becomes very small and PRESTO is able to recover the true partitions, i.c., St ~ S₃ (2) We observe that PRESTO finds it relatively difficult to correctly assign the instances on 5% and S to the right mixture components. This is due to two reasons: first. the instances on St and Sg fit naturally to nonlinear SVMs, whereas attempt to fit a linear SVM; second, the instances in S and S5. which are around (0, 0), are close to each other- this poses difficulty in demarcating their model components.
22

---

Under review as a conference paper at ICLR 2023
F ADDITIONAL EXPERIMENTS ON REAL DATA
Accuracy vs K. Here, we aim to assess the impact number of partitions (K) on the accuracy. K is varied from 6. Figure shows the variation of K with the test accuracy and the Macro-F1 score. We observe that the performance is stable across these range.
Perforaacc 7
Performance 82
Mivro P1
Arratary
Performance
70
Necminy Mario Fl
Performance
Accrracy
87
Macto
89.9
60
89.8
78
50
86
5
3
6
3
-+
K -
K
-
(a) CIFARIO
(b) PMNIST
DMNIST
(d) SVHN
Figure 7: Accuracy vs. the number of partitions K
Accuracy vs GPU usage for DMNIST and SVHN. Here, we perform the same experiments on DMNIST and SVHN as Figure 3
Figure 8 summarizes the results, which reveal same insights as
Figure 3
SVHN
76
DMNIST
91
Preeto
SVHN
76
DMNIST
90
90
Acquracy
89
Prosoo
80
KD
DCKD
DCKD
87
50
270
87
GPU (ln MLB)
GPU (In MIB)
GPU (In MIB)
GPU (ln MIB) -
(a) Training
(b) Inference
Figure& 8: Trade off between accuracy and maximum GPU memory for KD, DGKD and PRESTO during training (panel (a)) and inference (panel (b)) for DMNIST and SVHN.
G EXPLANATION OF PRESTO (ALGORITHM D
In PRESTO we aim to minimize monotone approximate-suulee subbmouua function G. In contrast to submodular or approximate submodular maximization where one can use greedy algorithms, min- imizing the approximate submodular function G needs completely different approach. Here, we resort to minimizing an upper bound m of G reducing the value of this upper bound m will ensure the value of the underlying function G remains low. Now, such an upper bound has to be such that the gap between m and G is low and maximizing m convenient. To this aim, we choose m to be modular. One can connect such modular approximation of a set function to simple linear approximation of a complex nonlinear function in the context of continuous optimization.
Now minimizing the modular function mg(S) wrt S is equivalent taking the minimum of [mg]( [(i, k)] over different (i, k) pairs ( since, modular(Set) = Eee modular(e)). In each we simply compute the minimum of m and update S accordingly. Note
iteration of Algorithm [ that:
m mf =(()
Σ
G(i,(i))) {(i,b)))+ + Σ agG((iAA)|5) ((),)))
(39)
(i,))8ns
A((
Σ
aG
(
Now, maximizing mg [i, k)) is same as maximizing og k) $ ((i,k)) or G((i.) ()
depending
ag
on whether k) E (3nd term in Eq 39 or k) & $ (4th term in Eq (39)). We compute these two quantities in L15 (line no. 15 in Algorithm and L.7 of the algorithm and store the values of
23

---

Under review as a conference paper at ICLR 2023
at different pairs (i, k) in the matrix M. Finally, take the minimum of M to find (i" k*) (L19) which indicates that the partition k* should contain 1". Therefore, include i" to Sy- (I.20). Now, since any instance can belong to exactly one partition, we remove i" from all other partition & + k* (L21). Finally, we update $ (L22).
H ADDITIONAL EXPERIMENTS WITH DATASETS HAVING LABEL NOISE
The datasets are homogenous, extremely balanced and there is lack of noise. In this section, present experiments where we added different amounts of label noise. Specifically, we changed each label to wrong label y with probability 10% and then probe the performance. The results for We observe that PRESTO performs much better
CIFAR10 10 and DMNIST presented in Table than the clustering baselines.
Accuracy P(ù y)
Method
CIFARTO
DMNIST
PRESTO
89.01
75.31
EqKMeans
88.49
74.46
Kmeans++
88.08
74.11
Agglomerative 87.59
74.46
Table 9: Comparison of classification accuracy P y y) of PRESTO against three unsupervised par- titioning methods (Equal- -means (Bennett et al. 2000). Kmeans++ (Arthur & Vassilvitskii [2006), Agglomerative (Müllner 2011), We change each label to a wrong label with probability 10%, thereby adding heterogenity to the dataset. Numbers in green and yellow indicate the best and the second best method.
I TIME COMPLEXITY ANALYSIS
Although PRESTO performs data partitioning while simultancously learning a set of mixture of models, the partitioning time (lines 19-22 in Algorithm I is negligible the bulk of the time is devoted to training (lines 10-18 in Algorithm M Since the training stage common to all the methods, PRESTO does not provide much disadvantage in terms of the time. Table T0 presents analysis of the per-iteration training time.
Training time per iteration
Method
CIFARTO
PMNIST
DMNIST
SVHN
PRESTO
9.6 min
16.6min
1.65 min
14.6 min
EqKMeans
9.4 min
16.6min
1.49 min
14.52 min
Kmeans++
7.96 min
16.6min
1.43 min
14.46 min
Table 10: Comparison of training time per iteration of PRESTO against Equal-Kmeant (Benneti al. 2000) and Kmeans++ (Arthur & Vassilvitskii 2006)
24


Reviews:
[{"review_id": "riMc1cmN32N", "reviewer": "Reviewer_jcju", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "", "summary_of_the_paper": "When the observed data comes from a population that consists of several subpopulations, it is useful to model each subpopulation separately. Hence, the paper proposes a method that simultaneously learns a partition of the dataset so that each subset of the data comes from the same distribution. The model parameters and the partitions are learned simultaneouly, which is formulated as an optimization problem. One major contribution of the paper is to design a numerical algorithm. ", "summary_of_the_review": "One major concern of the reviewer is that the difference of the proposed method and mixture of linear regression.  The paper is marginally below the acceptance threshold before the reviewer's comments are clarified.\n", "strength_and_weaknesses": "**Strengths**\nThe paper is well written and can be easily followed. \n\n**Weakness**\n1. One major concern of the reviewer is that the difference of the proposed method and mixture of linear regression is not highlighted.\n- In Section 2.3, the authors mentioned that mixture of linear regression is a special case of theirs. If so, then why you get different results for these two methods in Table 1. \n- Can the authors please clatify what's the difference between the proposed approach and the general EM algorithm used in fitting mixture of regression models.\n\n2. Numerical algorithm. An MM algorithm is proposed to minimize $G(\\mathbb{S})$. The success of MM algorithm depends heavily on the majorization function, which the reviewer believe shoudl be $m_{\\hat{\\mathbb{S}}}^{G}[\\mathbb{S}]$. However, the reviewer don't think $m_{\\hat{\\mathbb{S}}}^{G}[\\mathbb{S}]$ majorizes $G(\\mathbb{S})$ since the equality does not hold when $\\mathbb{S}=\\hat{\\mathbb{S}}$. Can the author please clarify?\n\n3. Experiment result. \n- Are the differences of different approaches in Table 1 significant?\n- The k-means approaches does not use any information about label $y$ but still performs as good as the proposed method, can the authors give some justificaitons on why this is the case?\n\n\n**Minor issues**\n\n1. The definitions of monotonity and $\\alpha$-submodular need to be double checked. The current version is not correct since the $f$ does not appear anywhere in the definition.\n2. The definition of relizable and non-relizable is not clear.\n3. The number of paritions is the same as the number of classes in the dataset, what happens if the number of partitions is overspecified or under specified?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "RoE8SzRwRGj", "reviewer": "Reviewer_f3vu", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper proposes a framework (called PRESTO) for simultaneously partitioning the data space and learning a separate model in each partition. The framework relies on solving an optimization problem with both discrete variables (data partitioning) and continuous parameters \n(model parameters) using a constrained optimization method.  ", "summary_of_the_review": "The paper proposes a framework (called PRESTO) for simultaneously partitioning the data space and learning a separate model in each partition. The framework relies on solving an optimization problem with both discrete variables (data partitioning) and continuous parameters (model parameters) using a constrained optimization method.\n\nThe idea is interesting and supported by theoretical results. \n\nHowever, there are several concerns regarding the manuscript:\n1) I think that PRESTO resembles with the Mixture of Experts (MoE) family of methods. Nevertheless, the paper does not include any reference and comparison with MoE.\n2) The authors compare only with the typical (trivial) case where partitioning is first obtained using clustering and then a separate model is trained on each partition. Even in this case, the results in Table 1 do not indicate any particular accuracy improvement.\n4) The time complexity of the method should be specified and compared to the methods in Table 1. It should be noted that the method must run for various values of K in order for the best number of partitions to be determined through cross-validation.\n5) I do not think that the method could be an efficient alternative for k-means of Bergman clustering as mentioned in section 2.3.\n6) If mixture of SVMs is a special case, why not implementing PRESTO with SVMs as classifiers and compare with the original mixture of SVMs method?\n7) In my opinion PRESTO does not perform mixture modeling. Mixture models are statistical models, PRESTO is a domain decomposition approach.\n8) The intuition behind the a_G hyperparameter should be provided.\n9) Does PRESTO algorithm converge? What is the role of the 'Iterations' parameter in Algorithm 1.  \n", "strength_and_weaknesses": "Stength\n - An interesting problem formulation is presented.\n - Partitioning is obtained through explicit specification of the partitions without resorting to a clustering model.\n - Theoretical bounds are provided.\n - An approximate solution to the constrained optimization problem is proposed.\n\nWeaknesses\n  - Method presentation (section 3.2) and algorithm description are not clear, thus they are difficult to comprehend.\n  - The method can not directly provide the partition assignments of new examples. Therefore an additional classification \n    model should be trained for the assignment task. \n  - Evaluation is insufficient: PRESTO is not compared to various methods belonging to the Mixture of Experts category.\n  - Time complexity is not reported", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "M4mbpE_KrFG", "reviewer": "Reviewer_dXtq", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper studies the following problem: Given a labeled dataset, a parameter $K$, and $K$ model architectures, we want to partition the data into $K$ partitions and also learn all the parameters for each model corresponding to the partitions. To help understand the problem and compare their approach, I will discuss a simple baseline they use in their paper.\n\nGiven a labeled dataset, first run $k$-means++ using only the features (ignoring the labels) to form the $K$ partitions, then independently train a model for each of the partitions (now also using the labels). Note that in this approach, we have completely ignored the labels in the data-partitioning step. What this paper does is to instead also incorporate the labels and simultaneously learn the model parameters and the data-partitions.", "summary_of_the_review": "Overall interesting paper. I am currently leaning to accept but would be willing to increase my score after the author response.", "strength_and_weaknesses": "Strengths: Interesting paper which is overall well-written and studies a significant problem. Such algorithms can be useful when there is a large dataset, which was potentially generated from heterogeneous data-sources.\n\nWeaknesses: \nAlthough the problem studied by the paper is interesting, the algorithm they provide (PRESTO) is a bit complicated. Some of the simpler baselines (like the k-means++ based approach) give performance which is pretty close to their approach. I am very skeptical if any practitioner would bother using their approach for the marginal improvement in utility compared to the simpler baselines. Also, can the authors comment on the time taken by their algorithm when compared to their baselines?\n\nAlthough the paper provides a approximation bound (Theorem 4), I am skeptical if this is of any use in practice. More concretely, for some of the loss functions and the data-sets which the authors consider in the experiments, what is the value of this approximation factor?\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject