Fair Model-based Clustering

AI summary

Overview

Research area: Fair machine learning for unsupervised learning, specifically fair clustering built on finite mixture models (stat.ML).

Technical level: Advanced. The paper assumes familiarity with finite mixture models, the Expectation-Maximization (EM) algorithm, gradient-based optimization, and common clustering fairness metrics (Balance, Gap).

Scope: The paper proposes Fair Model-based Clustering (FMC), a probabilistic clustering framework that enforces a fairness constraint on a finite mixture model, with two learning algorithms (FMC-GD and FMC-EM), theoretical guarantees for sub-sampled fairness constraints, and empirical evaluation on four UCI datasets including a 2.4-million-sample Census dataset.

What This Paper Is About

Fair clustering seeks clusters in which the proportion of sensitive attributes (such as gender or race) in each cluster mirrors the proportions in the full dataset. Most existing fair clustering methods extend K-means and must learn one cluster assignment per training point, so the number of learnable parameters grows with sample size, mini-batch learning is impossible, and assigning new unseen data fairly is hard. This paper instead places a fairness constraint directly on a finite mixture model, so the number of parameters depends only on the number of clusters and mixture parameters, not on the sample size.

Key Contributions

1. A new fair clustering algorithm based on a probabilistic mixture model, where the number of learnable parameters is independent of sample size, enabling scalable mini-batch learning and fair re-assignment of newly arrived (unseen) data.
2. Two learning procedures, FMC-GD (gradient descent) and FMC-EM (a Generalized EM variant), that optimize a log-likelihood objective penalized by a fairness term Δ(Θ), with FMC-EM shown empirically to give a better cost-fairness trade-off with smaller variance.
3. A sub-sampled fairness constraint with a theoretical guarantee (Proposition 1) bounding the deviation between Δ(Θ) computed on a sub-sample and the population Δ(Θ), plus an extension of Δ to multinary sensitive attributes (M > 2).
4. Empirical evaluation showing FMC is competitive with or better than existing fair clustering baselines (SFC, VFC, FCA) on three moderate-scale datasets and on a large-scale Census dataset, and that FMC can natively handle categorical data via a multinoulli mixture model.

Main Findings

• Fewer parameters, easier scaling: Because the (soft) assignment map is the parametric posterior ψ_k(x; Θ) rather than a per-point learned assignment, the number of learnable parameters does not grow with N. This is the mechanism that makes mini-batch learning and sub-sample learning possible, unlike existing methods where the assignment map must be learned for the entire dataset at once.
• FMC-EM vs. FMC-GD: On the Adult dataset, FMC-EM achieves a better cost-fairness trade-off with smaller variance than FMC-GD across five random initializations. FMC-EM is therefore used for the remaining experiments.
• FMC vs. SFC: FMC has lower Cost than SFC at a near-perfect fairness level. SFC operates only at (near-perfect) fairness and appears as a single point, whereas FMC lets the user control the fairness level; SFC learns centers and assignments sequentially while FMC learns them simultaneously.
• FMC vs. VFC: VFC is slightly superior to FMC on the cost-fairness trade-off because VFC minimizes Cost explicitly (FMC maximizes log-likelihood, only indirectly related to Cost). However, VFC fails to control the fairness level over the full range—for example, VFC does not provide a fair clustering with Δ < 0.06 on the Bank dataset—and VFC often fails to converge without L2 normalization, which FMC does not require.
• Large-scale performance: On the Census dataset (N = 2,458,285), FMC with mini-batch learning is competitive with VFC and better than SFC, while retaining the ability to control the fairness level across the entire range.
• Computational speed: Average computation time (seconds, standard deviation) over five random initials on Census: VFC 5053.0 (915.5); SFC 2218.2 (535.1); FMC with mini-batch at n = 0.1N 253.0 (12.1). The paper describes FMC as about 20x faster than VFC.
• Mini-batch vs. sub-sample learning: On Census, sub-sample learning at sizes 1%, 3%, 5%, and 10% gives Cost (×10^5) of 11.167, 10.988, 10.973, and 10.916; Δ of 0.005, 0.004, 0.003, and 0.003; Balance of 0.851, 0.894, 0.883, and 0.881; and relative times of 53.7%, 65.8%, 71.4%, and 77.2%. Mini-batch at 10% gives Cost 10.857, Δ 0.002, Balance 0.896, and 100.0% time. Sub-sample learning at 5% or larger performs similarly to mini-batch learning at lower computational cost.
• Theoretical guarantee: Proposition 1 states that, for Θ in a restricted parameter set defined by constants ξ, ζ, and ν, the sub-sampled Δ deviates from the full-data Δ by at most 4C√(d/n′) + 8√(2 log(8/δ)/n′) with probability at least 1 − δ, where n′ = min{n_1, n_2} and C = C(ξ, ζ, ν, x_max). This bound motivates using a fixed sub-sampled Δ during training.
• Categorical data: FMC applies to categorical data through a categorical (multinoulli) mixture model, and the paper reports that FMC performs well with categorical data, whereas baseline methods do not natively handle categorical variables and require nontrivial modifications.
• Multinary sensitive attributes: The paper proposes a modification of Δ (Eq. 14) that averages pairwise gaps Δ_{s1,s2} across sensitive groups, and reports that the modified FMC performs well on Bank and Credit datasets where the sensitive attribute has three categories.
• Robustness to K and covariance: FMC performs well regardless of the choice of the number of clusters K—negative log-likelihood decreases as K increases while Δ stays sufficiently low—and the paper also studies the choice of covariance structure.

Methodology in Plain English

The starting point is the standard finite mixture model, which assumes each data point comes from one of K component distributions, with a mixture weight π_k for each component. Instead of treating clustering as a geometric assignment problem (as K-means does), the method treats a point's cluster as a latent variable and uses the posterior probability of that latent variable as a soft assignment map.

To make clustering fair, the authors define a quantity Δ(Θ) that measures the largest gap—across all clusters—between the average posterior membership probability of one sensitive group and that of another. When Δ(Θ) is small, each cluster's membership reflects the overall group proportions. They then maximize the data log-likelihood subject to Δ(Θ) ≤ ε, or equivalently maximize log-likelihood minus λ·Δ(Θ) for a Lagrange multiplier λ. This replaces the per-datapoint assignment parameters of K-means-based fair clustering with a fixed set of model parameters.

Two optimizers are given. FMC-GD runs gradient descent directly on the penalized objective, reparameterizing the mixture weights through a softmax for numerical stability. FMC-EM embeds the penalty in the EM Q-function: the E-step computes expected complete-data log-likelihood minus λ·Δ(Θ), and the M-step takes gradient steps that increase this Q-function (a Generalized EM scheme, since no closed form exists).

For large datasets, the log-likelihood term can be evaluated on mini-batches, but Δ depends on the whole dataset. The authors approximate Δ using a fixed random sub-sample of size n, and Proposition 1 guarantees that the sub-sampled Δ is close to the full-data Δ for parameters in a bounded region. A second scalability trick is sub-sample learning: fit Θ on a random sub-sample, then assign the remaining points using the fitted parametric assignment map—which is possible precisely because FMC yields a functional form for the assignment rather than a table of per-point labels.

Why This Matters

Impact on research: Fair clustering has been dominated by K-means-based geometric approaches whose parameter count scales with the dataset, blocking mini-batch training and complicating inference on new data. By moving fairness constraints into a likelihood-based mixture model, this paper connects algorithmic fairness to the well-developed theory and tooling of model-based clustering (EM, Generalized EM, convergence results for mini-batch EM), and it supplies a finite-sample bound for the sub-sampled fairness penalty. It also opens fair clustering to categorical data, which distance-based K-means methods handle poorly.

Real-world applications (drawn from the paper's stated domains):

• Credit scoring, where decisions must not systematically disadvantage protected demographic groups.
• Criminal risk assessment, where the distribution of outcomes across groups is socially sensitive.
• College admissions, where clustering and grouping can influence downstream selection.
• General clustering deployments named in the paper: time series analysis, audio and language modeling, recommendation systems, and image clustering, any of which may need group-proportion guarantees.

Industry relevance: The mini-batch and sub-sample results speak directly to production-scale constraints. On the 2.4-million-row Census dataset, mini-batch FMC ran in roughly 253 seconds versus roughly 5053 seconds for VFC and roughly 2218 seconds for SFC. Sub-sample learning at 5% or more matched mini-batch quality at lower cost, and the parametric assignment map allows new incoming data to be assigned fairly without re-solving an assignment problem over the entire dataset—features that matter for streaming or continuously arriving data.

Future Directions

• Extending the theoretical analysis beyond the Gaussian mixture setting: the deviation bound in Proposition 1 is stated for Θ in a bounded region of Gaussian parameters, so analogous guarantees for categorical or other mixture families are not established in the provided content.
• Closing the cost-fairness gap with VFC: the paper notes VFC is slightly superior on Cost because it explicitly minimizes Cost, while FMC maximizes log-likelihood. Whether FMC can be adapted to target Cost directly is left open.
• Broader fairness notions: this work uses ratio-based Balance and difference-based Gap/Δ. Whether the framework extends to other fairness definitions for clustering is not addressed in the available content.
• Serving arriving data in a genuinely sequential manner: the paper argues the parametric assignment map makes fair assignment of new data easy and mentions the naive fixed-center solution is inadequate when data arrive sequentially, but a fully sequential procedure is not developed in the visible text.

Target Audience

Researchers and graduate students in machine learning, statistics, and algorithmic fairness who are familiar with mixture models and EM, and who want a scalable, likelihood-based alternative to K-means-derived fair clustering. It is also relevant to applied practitioners working with large tabular datasets that mix continuous and categorical features and that carry group-proportion requirements, such as teams building credit, admissions, or risk-assessment systems. Readers without background in mixture models or EM will need to consult the referenced literature first.

Authors’ abstract

The goal of fair clustering is to find clusters such that the proportion of sensitive attributes (e.g., gender, race, etc.) in each cluster is similar to that of the entire dataset. Various fair clustering algorithms have been proposed that modify standard K-means clustering to satisfy a given fairness constraint. A critical limitation of several existing fair clustering algorithms is that the number of parameters to be learned is proportional to the sample size because the cluster assignment of each datum should be optimized simultaneously with the cluster center, and thus scaling up the algorithms is difficult. In this paper, we propose a new fair clustering algorithm based on a finite mixture model, called Fair Model-based Clustering (FMC). A main advantage of FMC is that the number of learnable parameters is independent of the sample size and thus can be scaled up easily. In particular, mini-batch learning is possible to obtain clusters that are approximately fair. Moreover, FMC can be applied to non-metric data (e.g., categorical data) as long as the likelihood is well-defined. Theoretical and empirical justifications for the superiority of the proposed algorithm are provided.