On Distributed Larger-Than-Memory Subset Selection with Pairwise Submodular Functions

On Distributed Larger-Than-Memory Subset Selection with Pairwise Submodular Functions

Speakers: Maximilian Böther, Research Scientist; Abraham Sebastian, Research Scientist; Pranjal Awasthi, Research Scientist; Ana Klimovic, Assistant Professor; Srikumar Ramalingam, Research Scientist

Conference: MLSys 2025

YouTube: https://slideslive.com/39043019

Overview

In the realm of large-scale machine learning, the cost associated with training models on massive, often petabyte-scale datasets presents a significant challenge. This talk, presented by Maximilian Böther as a joint work from Google and ETH, introduces a novel approach to distributed larger-than-memory subset selection using pairwise submodular functions. The core problem addressed is the need to efficiently identify a small, representative subset of a colossal dataset that can serve as a proxy for the full dataset during model training and experimentation. Such a subset aims to provide an accurate indication of a model's performance on the entire dataset, but at a fraction of the computational cost.

The fundamental issue preventing direct application of existing subset selection algorithms is that even the selected subset often exceeds the DRAM capacity of a single machine. Traditional algorithms, which typically assume a global view of the data and in-memory processing, are rendered ineffective in such "larger-than-memory" scenarios. Böther and his team tackle this by proposing a distributed selection procedure that not only handles the scale but also strives to maintain the performance characteristics achievable by centralized, global-view algorithms. Their work is particularly relevant for large-scale ML organizations like Google, where engineers frequently run costly experiments on constantly evolving, massive datasets.

The talk delves into the theoretical underpinnings of submodularity and pairwise submodularity, which are crucial for defining what constitutes a "good" subset by balancing data utility and diversity. It then progressively introduces distributed strategies, starting with a naive partitioning approach that highlights inherent performance trade-offs, and culminating in an advanced preprocessing technique called bounding. This bounding mechanism is designed to provably include or exclude data points, enabling arbitrary parallelization without the performance degradation typically associated with distributed greedy submodular optimization. The presented empirical results demonstrate that this approach not only scales effectively but can also yield subsets that perform comparably to, or even better than, those selected by centralized greedy methods, offering substantial compute savings for repeated experiments.

Background

▶ Watch: Introduction: Larger-than-memory subset selection problem (0:00)

The explosion of data in modern machine learning has led to datasets that routinely span petabytes, making model training and iterative experimentation extremely resource-intensive. Companies operating at this scale face immense computational costs, driving the need for efficient data management strategies. One such strategy is subset selection, where the goal is to identify a smaller, yet highly representative, portion of the full dataset. This smaller subset can then be used for faster model prototyping, hyperparameter tuning, or even full training, with the expectation that insights gained from the subset generalize well to the full dataset.

The primary hurdle for applying traditional subset selection algorithms to these massive datasets is the memory constraint. Most existing algorithms, particularly those based on greedy optimization, assume that the entire dataset or at least the actively selected subset can reside in the DRAM of a single machine. This assumption breaks down when dealing with datasets so large that even a small fraction of them still exceeds single-machine memory limits, leading to what the authors term "larger-than-memory" subset selection. This necessitates a distributed approach.

The theoretical foundation for this work lies in submodularity, a property of set functions that is often described as the discrete analog of convexity. A key characteristic of submodular functions is the principle of diminishing returns: the marginal gain of adding an element to a set decreases as the set itself grows larger. Böther illustrates this with an example of counting distinct colors in an urn: adding a new ball is more likely to introduce a new color (and thus yield a higher "gain") if the urn already contains fewer distinct colors. Formally, for a set function F and sets A ⊆ B, the gain of adding an element x to A (F(A ∪ {x}) - F(A)) is greater than or equal to the gain of adding x to B (F(B ∪ {x}) - F(B)). This property is highly desirable for subset selection, as it encourages diversity and prevents over-representation of similar items.

To define a "good" subset, the talk focuses on pairwise submodular functions, which combine two critical aspects: utility and diversity.

• Utility quantifies how "good" an individual data point is. This can be measured using various machine learning-specific metrics, such as a model's loss on that data point, its margin in a classification task, or other uncertainty measures. The algorithm is designed to be agnostic to the specific utility metric chosen.
• Diversity captures how distinct data points are from each other. This is typically modeled using similarity metrics (e.g., cosine similarity) and relationships derived from k-nearest neighbor (k-NN) graphs. Computing k-NN graphs at scale is a non-trivial task, but Böther notes that Google's extensive infrastructure for tasks like advertisement targeting can be leveraged for efficient similarity computations.

While higher-order interactions (considering three or more data points simultaneously) could theoretically offer a more nuanced view of a subset, they are currently computationally infeasible at the scales considered. Therefore, pairwise submodularity strikes a practical balance, capturing essential aspects of data quality and variety within computational limits.

The standard approach for submodular subset selection in a centralized setting is a greedy algorithm. This iterative procedure starts with an empty set and, in each step, adds the data point that maximizes the marginal gain to the current subset. This continues until the desired subset size k is reached. While efficient implementations exist (e.g., using heap data structures or priority queues to quickly find the element with maximum marginal gain), this algorithm inherently requires a global view of the data—specifically, the current subset and all relevant similarity values—which makes it unsuitable for larger-than-memory scenarios. The challenge, therefore, is to distribute this greedy selection process effectively without sacrificing the quality of the selected subset.

Key Findings

▶ Watch: Understanding submodularity: The diminishing returns property (0:50)

The central findings of this research revolve around enabling efficient and high-quality submodular subset selection at petabyte scales, overcoming the limitations of single-machine memory.

Firstly, the talk highlights the inherent trade-off in a naive distributed partitioning approach. While distributing the data and running local greedy algorithms on partitions allows for parallel computation, empirical results demonstrate a clear degradation in the quality of the selected subset as the number of partitions (and thus parallelization) increases. Conversely, increasing the number of rounds in this iterative partitioning scheme can mitigate this error, but at the cost of increased computation. This establishes a critical problem: how to parallelize submodular optimization without compromising performance.

Secondly, the introduction of bounding as a preprocessing step is a major contribution. Bounding is an iterative procedure designed to provably identify data points that will definitely be part of the optimal subset and those that will definitely not. This allows for the reduction of the ground set before running any greedy selection. Crucially, the bounding procedure itself can be arbitrarily parallelized using distributed data processing frameworks (like Google's Beam or Flume) without incurring the performance penalty observed in the naive partitioning method. This solves the core problem of scalable, performant subset selection.

Thirdly, the experimental evaluation of bounding reveals its significant effectiveness. In some scenarios, exact bounding (without sampling) not only matches but can even slightly exceed the performance of the centralized greedy algorithm. This is attributed to the fact that the centralized greedy algorithm is itself a heuristic and not guaranteed to find the global optimum. The ability to achieve scores above 100% (relative to the centralized greedy baseline) underscores the power of bounding in identifying high-quality subsets.

Finally, approximate bounding (using sampling to estimate neighborhood utility) is shown to be highly practical. For instance, with just 30% sampling of neighbors, the algorithm was able to exclude 50% of the dataset from consideration before running the greedy procedure. This translates directly into substantial computational savings, as the subsequent greedy selection needs to operate on a significantly smaller ground set. These findings collectively demonstrate a robust and scalable solution for larger-than-memory submodular subset selection, offering a pathway to dramatically reduce the cost of ML experimentation on massive datasets.

Technical Deep Dive

▶ Watch: Pairwise submodularity: Modeling utility and diversity (2:00)

The technical core of this work addresses the challenge of scaling submodular subset selection beyond single-machine memory limits. The talk outlines two main approaches: an initial, simpler distributed strategy and a more sophisticated "bounding" technique.

Naive Distributed Partitioning and Iterative Refinement

The first attempt at distribution involves a straightforward partitioning strategy. The entire ground set G of data points is divided into m disjoint subsets, each assigned to a separate machine. Each of these m machines then independently executes the standard centralized greedy algorithm on its local partition. This local greedy selection identifies the best data points within that specific partition.

The results from all m machines are then union-ed. A critical aspect here is that this union operation does not require materializing the entire combined result in a single machine's memory. Instead, it leverages distributed data processing frameworks (like Apache Beam or Google's internal Flume) that support out-of-memory compute and can handle data distributed across a file system.

To improve the quality of the selected subset, this partitioning and unioning process can be repeated over multiple rounds. In each subsequent round, the unioned result from the previous round becomes the new ground set, which is then re-partitioned (often randomly) and processed again.

• The target subset size k can be adjusted per round, typically shrinking linearly towards the final desired size.
• The number of machines used can also be reduced in later rounds to minimize resource utilization and potential "rounding error" (the impact of local greedy choices on global optimality).

Visual Example: Imagine a dataset with data point IDs (1-9).

1. Round 1: Data is partitioned (e.g., {1,2,3}, {4,5,6}, {7,8,9}) across three machines. Each machine selects, say, two best local points. Results are unioned (e.g., {2,3,5,6,8,9}).
2. Round 2: The unioned set is re-partitioned (e.g., {2,5,8}, {3,6,9}). Each machine selects one best local point. Results are unioned (e.g., {2,6,9}) as the final subset.

A key practical distinction from prior theoretical work is that this approach does not re-run the greedy algorithm on the final unioned result. While such a step is often required for certain mathematical guarantees, the authors prioritize practical computability and empirically demonstrate its performance.

Limitations: Empirical evaluation on CIFAR showed that while more rounds improved the objective score, more partitions (increased parallelization) decreased the objective score. This highlights a fundamental trade-off: naive distribution sacrifices global optimality for parallelization, indicating the need for a more sophisticated approach.

Bounding: A Performance-Preserving Distributed Preprocessing Step

To overcome the performance degradation of naive parallelization, the authors introduce bounding. This is an iterative preprocessing procedure that runs until convergence, designed to identify data points that are either provably included or provably excluded from the final optimal subset. After bounding, the remaining, reduced set can then be fed into the greedy submodular routine. The critical advantage of bounding is its ability to be arbitrarily parallelized using distributed data processing frameworks (like Beam/Flume) without hurting the performance of the subset selection.

The core idea of bounding relies on comparing minimum and maximum utility values for data points within the context of a potential solution.

• Maximum Utility: For a given data point, its maximum utility is considered by only taking into account its neighbors that are already part of the current partial solution. The assumption is that if we pick other similar neighbors, the overall score of the solution would decrease due to diminishing returns (increased similarity).
• Minimum Utility: This generally refers to the worst-case utility a point could provide, often assuming all its neighbors are already included in the set (which would minimize its marginal gain).

The bounding procedure consists of two main steps, repeated iteratively:

1. Shrink Step: The goal is to discard points that will provably never be included in the final subset.

• Data points are sorted by their minimum utility and maximum utility.
• A threshold is established (e.g., based on the minimum utility of the k-th best point if we were to select k items).
• If a data point's maximum utility is smaller than the minimum utility of a high-ranking (e.g., top k) data point, that point can be provably removed from the ground set. The intuition is that no matter what, the high-ranking point will always offer a better gain.
• Example: If point A's max utility is 5, and point B's min utility is 7, point A will never be preferred over point B.

1. Growth Step: The goal is to include points that will provably always be part of the final subset.

• Again, data points are sorted by maximum and minimum utility.
• If a data point's minimum utility is higher than the maximum utility of a low-ranking data point (e.g., a point just outside the top k), then this data point can be provably included in the solution.
• Example: If point C's min utility is 10, and point D's max utility is 8, point C will always be preferred over point D and thus included.

These steps are repeated, updating the ground set and partial solution, until no more points can be provably included or excluded (i.e., the procedure converges).

Distributed Implementation Challenges: A significant challenge in distributed bounding is efficiently checking neighbor relationships. In a centralized setting, this is an O(1) hash table lookup. In a distributed environment, determining if a data point is a neighbor of another requires complex operations, often involving three-way joins within data processing frameworks like Beam. The paper provides implementation details on how to manage these distributed checks efficiently.

Approximate Bounding

To further enhance practicality, the concept of approximate bounding is introduced. Instead of computing the exact minimum utility (which assumes a worst-case scenario where all neighbors of a point are already used), approximate bounding involves sampling a neighborhood of a data point. This allows for the calculation of an expected utility rather than a strict minimum. This approximation reduces computational overhead while still yielding significant benefits in terms of dataset reduction.

By leveraging distributed data processing frameworks for bounding, the system can scale across an arbitrary number of machines without the performance degradation observed with naive parallelization. This ensures that the quality of the selected subset remains high, even for petabyte-scale datasets.

Experimental Setup & Results

▶ Watch: Centralized greedy algorithm for subset selection (3:15)

The experimental evaluation primarily focused on demonstrating the effectiveness of the proposed distributed subset selection methods, particularly bounding, against a centralized greedy baseline.

Initial Evaluation (CIFAR Dataset):

For the initial distributed partitioning approach, experiments were conducted on the CIFAR dataset. While CIFAR is not a petabyte-scale dataset, it allowed for direct comparison with a centralized greedy algorithm to understand the performance characteristics of the distributed methods.

• Utility Metric: Margin-based uncertainty was used, which is common in classification tasks (e.g., for a neural network, a data point with a smaller margin between its predicted class probability and the next highest probability is considered more "uncertain" and potentially more informative).
• Similarity Metric: Cosine similarity was employed, computed based on features extracted from the data points.
• Diversity Modeling: A 10-nearest neighbor (k-NN) graph was constructed to capture local similarities.
• Objective Score: Results were normalized, with the centralized greedy algorithm's score set as 100 for comparison.
• Hyperparameters: The number of partitions (m) and the number of rounds were varied.

Results of Naive Partitioning:

The heatmap presented in the talk clearly illustrated the trade-offs:

• More rounds (Y-axis) generally increased the objective score, meaning the quality of the selected subset improved. This is intuitive, as more intermediate steps reduce the "rounding error" of local greedy decisions.
• More partitions (X-axis) generally decreased the objective score. This was the critical finding: simply parallelizing the greedy algorithm by distributing it across more machines (more partitions) led to a worse-performing subset. This highlighted the need for a more sophisticated distribution strategy that doesn't sacrifice performance for parallelism.

Evaluation of Bounding:

The bounding algorithm was evaluated under different configurations:

1. No Sampling (Exact Bounding): This corresponds to the exact bounding procedure, where minimum and maximum utilities are computed precisely (e.g., considering all neighbors for minimum utility).
2. 30% Sampling: Approximate bounding where only 30% of a data point's neighbors are sampled to estimate its expected utility.
3. 70% Sampling: Approximate bounding where 70% of a data point's neighbors are sampled.

Bounding Results:

The results for bounding were highly positive:

• Performance Relative to Centralized: In several scenarios, the exact bounding algorithm (no sampling) achieved objective scores above 100% (e.g., 100.08, 100.12). This is a significant finding, as it implies that the distributed bounding approach can sometimes yield a better subset than the centralized greedy algorithm itself, which is a heuristic and not guaranteed to find the global optimum.
• Dataset Reduction: A particularly impactful result was from the 30% sampling scenario, where the bounding procedure was able to provably exclude 50% of the dataset before the greedy selection even began. This demonstrates immense potential for computational savings. The number of points provably included was also reported (e.g., 1.5% for 30% sampling), indicating that a significant portion of the decision-making can be offloaded to the highly parallelizable bounding step.
• Scalability: The key takeaway from the bounding experiments was that it performs as well as or better than the centralized algorithm in many scenarios, and crucially, it is scalable. The authors emphasized that "we can just throw a huge cluster at it and we don't really have to wait and we don't really suffer if we use more machines." This directly addresses the performance degradation issue observed with naive parallelization.

Generalizability and Downstream Performance:

The talk also touched upon the generalizability of the approach and potential downstream impacts:

• Modality Agnostic: The algorithm is agnostic to the data modality (images, text, etc.) as long as meaningful utility and similarity values can be computed. For text data, utility could be derived from loss values per token sequence or gradient norms, and similarity metrics (an active research area) can be employed.
• Utility Computation: The utility values (e.g., loss or margin) typically require a forward pass on a pre-trained model. This initial cost is justified for datasets that are reused multiple times (e.g., ImageNet for many experiments), where the upfront computation pays off in subsequent savings.
• Downstream Performance Gap: Böther candidly acknowledged that there definitely is a gap in downstream task performance when training on a subset compared to the full dataset. This gap is highly dependent on hyperparameters (e.g., alpha for utility-diversity balance, specific utility/similarity metrics). However, for a 50% subset, performance drops of only 2-3% have been observed for classification tasks using margin-based utility, which is often an acceptable trade-off for the substantial compute savings. The potential is even greater for noisy datasets like those used for large language models (LLMs), where subset selection could significantly reduce training costs.

In summary, the experimental results clearly validate bounding as an effective, scalable, and performance-preserving method for distributed larger-than-memory subset selection, offering substantial resource efficiency for ML practitioners.

Practical Implications

▶ Watch: Distributed approach: Partitioning data and iterative rounds (4:10)

The research on distributed larger-than-memory subset selection carries significant practical implications for various stakeholders in the machine learning ecosystem, from individual practitioners to large infrastructure teams.

For ML Practitioners and Researchers:

• Accelerated Experimentation: The most immediate benefit is the ability to conduct experiments much faster. By training models on smaller, representative subsets, researchers can iterate on model architectures, hyperparameter choices, and feature engineering at a significantly accelerated pace, reducing the time-to-insight.
• Cost Savings: For companies operating at petabyte scales, the computational cost of training models is immense. This work offers a direct path to substantial resource savings. As demonstrated by the ability to exclude 50% of a dataset, this translates to potentially halving GPU utilization and associated energy costs for repeated experiments on large, static datasets like ImageNet. This is particularly relevant for scenarios where utility values (e.g., forward pass on a pre-trained model) can be computed once and reused for many subsequent experiments.
• Handling Massive Datasets: Practitioners are no longer constrained by the memory limits of a single machine when selecting representative subsets. This opens up possibilities for working with datasets that were previously intractable for such methods.
• Focus on Core ML Problems: By automating and scaling the data selection process, ML engineers can dedicate more time to model development and problem-solving, rather than wrestling with data infrastructure challenges.

For Infrastructure Teams:

• Leveraging Existing Distributed Systems: The proposed solution is designed to integrate with existing distributed data processing frameworks like Apache Beam or Google's internal Flume. This means infrastructure teams can leverage their established tooling and expertise, rather than needing to build entirely new systems.
• Efficient Resource Utilization: The bounding step, being arbitrarily parallelizable, allows for efficient utilization of large clusters. It enables scaling out computation without diminishing the quality of the output, ensuring that invested hardware resources are effectively used.
• Out-of-Memory Processing: The architecture inherently supports out-of-memory processing, alleviating the burden on infra teams to manage complex in-memory caching or data sharding strategies for large-scale data.

For Model Builders and Deployers:

• Faster Model Prototyping: The ability to quickly generate high-quality subsets means model builders can rapidly prototype new models or adapt existing ones to new tasks, reducing development cycles.
• Data Quality Control: While not explicitly a quality control tool, the principles of utility and diversity can implicitly guide data curation efforts by highlighting important or unique data points.
• Informed Trade-offs: Model deployers can make informed decisions about the trade-off between model performance and training cost. If a 2-3% performance drop is acceptable for a 50% reduction in training data, this provides a clear pathway to more efficient deployment.

Tradeoffs and Limitations:

• Initial Utility Calculation Cost: The method requires an initial computation of utility values for each data point (e.g., a forward pass on a pre-trained model). This upfront cost makes the approach less suitable for "one-off" datasets where the cost of utility calculation might outweigh the savings from subset selection. It is most beneficial for frequently reused datasets.
• Hyperparameter Tuning: The quality of the selected subset and its impact on downstream model performance are highly dependent on the choice of hyperparameters, specifically the balance between utility and diversity (e.g., alpha), and the specific utility and similarity metrics used. Careful tuning is required to achieve optimal results for a given task.
• Potential Performance Gap: While often small and acceptable (e.g., 2-3% accuracy loss for 50% subset), there is an acknowledged performance gap between models trained on subsets and those trained on full datasets. Practitioners must evaluate if the resource savings justify this potential reduction in peak performance for their specific application.
• Agnosticism to Metrics: While the algorithm is agnostic to the specific utility and similarity metrics, choosing effective metrics is crucial. Poorly chosen metrics will lead to a sub-optimal subset.
• Computational Complexity of Distributed Joins: While bounding is parallelizable, the underlying distributed operations, particularly multi-way joins for neighbor checks, can be computationally intensive and require careful optimization within the distributed data processing framework.
• No Higher-Order Interactions: The pairwise submodular function formulation, while practical, does not capture higher-order interactions between three or more data points. This is a current algorithmic limitation for scalability.

Despite these considerations, the presented distributed subset selection framework provides a powerful tool for navigating the challenges of large-scale machine learning, enabling more efficient and cost-effective development cycles.

Key Takeaways

• Distributed Subset Selection is Essential for Scale: Traditional subset selection algorithms fail when datasets, or even the resulting subsets, exceed single-machine memory, necessitating distributed approaches for petabyte-scale ML.
• Naive Parallelization Degrades Performance: Simply partitioning data and running local greedy submodular optimization leads to a decrease in subset quality as parallelization increases, highlighting a critical trade-off.
• "Bounding" Enables Scalable, High-Quality Selection: The proposed iterative "bounding" preprocessing step can be arbitrarily parallelized without performance degradation, effectively identifying data points that are provably included or excluded from the optimal subset.
• Bounding Can Outperform Centralized Greedy: Empirical results show that bounding, particularly exact bounding, can achieve objective scores slightly higher than the centralized greedy algorithm, which is itself a heuristic.
• Significant Resource Savings Achieved: Approximate bounding (e.g., with 30% neighbor sampling) can exclude up to 50% of the dataset prior to greedy selection, leading to substantial computational savings for future model training runs.
• Modality-Agnostic and Reusable: The framework is applicable across different data modalities (images, text) as long as utility and similarity metrics can be defined, and is most beneficial for frequently reused large datasets where initial utility computation costs are amortized.
• Trade-offs Exist, but are Manageable: While a small performance gap (e.g., 2-3% accuracy loss for a 50% subset) on downstream tasks may occur, it is often an acceptable trade-off for the significant resource savings, especially with careful hyperparameter tuning.

About the Speaker(s)

The talk was presented by Maximilian Böther, who is affiliated with both Google and ETH Zurich. His work on distributed larger-than-memory subset selection is a direct result of an internship he conducted at Google's BigML group. He collaborated with a team of researchers from Google and ETH, including Abraham Sebastian, Pranjal Awasthi, Ana Klimovic, and Srikumar Ramalingam. Srikumar Ramalingam, one of the co-authors, was specifically mentioned in the Q&A session regarding his work on similarity metrics for text data, indicating the team's broader expertise in related ML challenges. The collective expertise of this group from leading industry and academic institutions underpins the robust and practical nature of the research presented.