Traditional ML often focuses on training a single model to minimize the average loss across all data points. However, this approach may not be optimal when dealing with heterogeneous data or when the data comes from multiple sources. We propose a novel "sum-of-minimum" model by considering an ensemble of $k$ models and optimizing them jointly. The key idea is to find the optimal assignment of each data point to the model that performs best on it, while simultaneously optimizing the performances of these models. This approach is mathematically formulated as:
minimize_{x_1,...,x_k} sum_{i=1}^n min{loss(x_1, data_i), ..., loss(x_k, data_i)}
where x_1,...,x_k represent the parameters of the k models. By minimizing the sum of minimum losses, the model encourages specialization among the k models. Each model is incentivized to specialize in a subset of the data points where it achieves the lowest loss.
However, solving this optimization problem is challenging due to the non-smoothness of the minimum function and the non-convexity of the objective. We present an algorithm that approximately solves the problem, consisting of an initialization step motivated by k-means++ and a sequence of iterations similar to Lloyd's algorithm. Under certain assumptions, we provide tight performance and convergence bounds for our algorithm.
We demonstrate the effectiveness of our approach through experiments on various tasks, including generalized principal component analysis, neural network training, and mixed linear regression. We believe that the "sum-of-minimum" model can significantly improve ML performance when dealing with diverse and complex datasets.