Disentangling Marginal Capacity from Geometry in Optimal-Transport Label Transfer
Listen
Radio episode about this paper
Transcript
Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.
Tom: Today's paper: "Disentangling Marginal Capacity from Geometry in Optimal-Transport Label Transfer".
Jane: The gist The source marginal imposes a class-level transport capacity, so a class can be starved when its target demand exceeds its available mass,
Tom: First, who's behind it and why it matters.
Paper summary: Tom: So we’re looking at this paper called "Disentangling Marginal Capacity from Geometry in Optimal-Transport Label Transfer." Basically, the core idea is that optimal transport, or OT, which we use to align distributions and transfer labels between source and target data, has a problem with rare classes.
Jane: Right. The paper identifies a specific source of failure in the transport marginal itself. It claims that when we use the standard uniform weighting for the source points, each class gets a transport capacity proportional to its frequency in the source data.
Tom: That means if a class is rare, it might not have enough mass to satisfy its target demand even if the embedding space separates it perfectly well from other classes.
Lu: The paper shows that this displaced mass doesn't just get lost; instead, it gets preferentially absorbed by the high-capacity classes in the target distribution rather than by classes that are geometrically similar in the embedding.
Meng: So, we see these confident errors because of this capacity mismatch, which is hard to spot just by looking at the transport plan itself.
Jane: The results show that when you look at read-outs like entropic smoothing or marginal relaxation, they don't recover much of that missing performance.
Tom: But the authors test different settings for parameters like entropic regularization epsilon, marginal penalty tau, and source marginal mu to see what helps.
Lu: They find that matching the transport marginal to the target composition substantially improves transfer, but only for those classes that were capacity-limited in the first place.
Jane: So the intervention isn't just tweaking how we regularize; it’s about assigning mass directly to the classes that need it, instead of changing things like read-outs.
Tom: The final proposed correction is to assign source mass in proportion to the target composition rather than based on the source frequency.
Jane: And when they do this, they see a macro F1 score increase from zero point three two two up to zero point six two three on single-cell data, which is a big jump <ref:2610.14853#pg2>.
Lu: What’s important here is that correcting the marginal removes this capacity-driven component of error, and what’s left over reflects the actual limitations of the representation geometry itself.
Meng: So it’s not just a representation problem; it’s also about how much mass we are actually putting into each class budget.
Conclusion: Tom: So, thinking about "Disentangling Marginal Capacity from Geometry in Optimal-Transport Label Transfer," it seems the main point is that transfer error isn't just about how well the embedding spaces are structured geometrically.
Jane: It’s about a separate issue, which they call marginal capacity. This capacity is determined by how much mass a source class has available to move, which isn't always proportional to its presence in the source data.
Lu: The paper suggests we have two distinct sources of error: one related to what’s missing in the transport plan because of insufficient budget, and another related to how well the representation captures true separability.
Meng: So if you only fix the geometry—like improving your embedding—you might still have problems if you haven't allocated enough source mass to cover those rare classes.
Jane: Exactly. The authors show that matching the transport marginal to the target composition fixes that capacity issue, revealing what’s left over as purely geometric error related to class separability.
Tom: The implication for us is that when we see poor label transfer, we should check both things: is the geometry bad, or are we starving a class because its budget isn't set correctly?
Lu: It tells us that fixing one part doesn't solve everything; you need to address the capacity constraint before you can fully assess the quality of your representation.
Jane: So, it’s a call to action for how we set up our OT problems: don't just focus on the distance between points; make sure your class budgets are realistic based on where you want them to end up.
Ali Khodabandeh Yalabadi, Azim Dehghani Amirabad, Artem Moskalev, Yun-Ching Chen, Mehdi Pirooznia, Scott Oloff, Rui Liao
Johnson & Johnson Innovative Medicine
cs.LG
Submitted: 2026-10-10
Updated: 2026-10-10
The gist: The gist The source marginal imposes a class-level transport capacity, so a class can be starved when its target demand exceeds its available mass, even when the embedding separates it well
Key concepts
- Marginal Capacity
- This is a constraint imposed by the source distribution on how much mass a specific class can transport during label transfer. If the target demand exceeds this limited capacity, the class can be starved of necessary mass, leading to failure regardless of how well classes are positioned in the embedding space.
- Capacity-driven Failure
- When capacity limits cause errors, displaced mass moves to classes with higher transport capacity rather than those that are close in the embedding. This is hard to detect because a class below capacity might have low recall but high precision, mimicking confusion with another class.
- Correcting the Marginal
- Instead of adjusting regularization parameters, the paper suggests assigning source mass proportionally to target composition. This correction removes the coverage-dependent error caused by capacity mismatch, leaving behind errors that reflect true class separability in the data representation.
Summary
The gist The source marginal imposes a class-level transport capacity, so a class can be starved when its target demand exceeds its available mass, even when the embedding separates it well
Capacity-driven failure
The paper identifies "marginal capacity as a source of rare-class transfer errors, with displaced mass flowing to high-capacity rather than geometrically similar classes This failure is difficult to detect because rare classes can be displaced into large, well-supported classes with high transport confidence The capacity signature shows that a class below capacity has low recall but high precision, while the largest class has high recall but low precision This is consistent with a class being short of transport mass rather than being confused with another class
Where the displaced mass goes
The analysis compares two explanations for the failure: limited representation geometry and limited transport capacity The results show that displaced mass follows available capacity rather than embedding proximity This is consistent with the categorical analysis where in 30 of 33 below-capacity observations, the absorbing class is the largest source class, while it is never the nearest class
Detection and interventions
The paper tests which components of the OT formulation can recover the missing rare-class performance by varying parameters like entropic regularization epsilon, marginal penalty tau, and source marginal mu Matching the marginal to the target composition substantially improves transfer, but only for capacity-limited classes The relevant intervention is to assign mass to the classes that need it, rather than to change the read-out or relax the OT regularization
Correcting the marginal
To correct the capacity mismatch, source mass is assigned in proportion to target composition rather than source frequency This correction substantially improves label transfer, with macro-F1 increasing from 0.322 to 0.623 on single-cell data Correcting the marginal removes a coverage-dependent component of label-transfer error and exposes a residual component associated with class separability
Capacity–geometry separation
Correcting the marginal removes the capacity constraint, so the remaining error should reflect limitations of the representation The residual component is associated with separability, not coverage This double dissociation shows that correcting the marginal removes a coverage-dependent component of transfer error while leaving a residual component associated with class separability
Estimating the target composition
The composition-matched marginal requires the target composition pi, which is unknown in practice The kNN pre-pass estimator closes 69% of the oracle gap on single-cell data, while JCPOT closes 59% on Office-Home The best deployable estimator is the kNN pre-pass there
Conclusion
The paper concludes that marginal capacity is a distinct source of transfer error, not simply another manifestation of poor representation quality Matching the marginal to the target composition removes this capacity-driven component and reveals the residual geometric error, but the correction is useful primarily where coverage is limiting and requires an estimate of the target composition The effective intervention is instead to change the class-level mass allocation itself Matching the marginal to the target composition removes this capacity-driven component and reveals the residual geometric error, but the correction is useful primarily where coverage is limiting and requires an estimate of the target composition The source marginal imposes a class-level transport capacity, so a class can be starved when its target demand exceeds its available mass, even when the embedding separates it well The displaced mass is absorbed by classes with greater capacity rather than by geometrically similar classes Correcting the marginal removes this coverage-dependent component of error, while the residual tracks class separability The effective intervention is instead to change the class-level mass allocation itself Matching the marginal to the target composition removes this capacity-driven component and reveals the residual geometric error, but the correction is useful primarily where coverage is limiting and requires an estimate of the target composition The source marginal imposes a class-level transport capacity, so a class can be starved when its target demand exceeds its available mass, even when the embedding separates it well The displaced mass is absorbed by classes with greater capacity rather than by geometrically similar classes Correcting the marginal removes this coverage-dependent component of error, while the residual tracks class separability The effective intervention is instead to change the class-level mass allocation itself Matching the marginal to the target composition removes this capacity-driven component and reveals the residual geometric error, but the correction is useful primarily where coverage is limiting and requires an estimate of the target composition The source marginal imposes a class-level transport capacity, so a class can be starved when its target demand exceeds its available mass, even when the embedding separates it well The displaced mass is absorbed by classes with greater capacity rather than by geometrically similar classes Correcting the marginal removes this coverage-dependent component of error, while the residual tracks class separability The effective intervention is instead to change the class-level mass allocation itself Matching the marginal to the target composition removes this capacity-driven component and reveals the residual geometric error, but the correction is useful primarily where coverage is limiting and requires an estimate of the target composition The source marginal imposes a class-level transport capacity, so a class can be starved when its target demand exceeds its available mass, even when the embedding separates it well The displaced mass is absorbed by classes with greater capacity rather than by geometrically similar classes Correcting the marginal removes this coverage-dependent component of error, while the residual tracks class separability The effective intervention is instead to change the class-level mass allocation itself Matching the marginal to the target composition removes this capacity-driven component and reveals the residual geometric error, but the correction is useful primarily where coverage is limiting and requires an estimate of the target composition The source marginal imposes a class-level transport capacity, so a class can be starved when its target demand exceeds its available mass, even when the embedding separates it well The displaced mass is absorbed by classes with greater capacity rather than by geometrically similar classes Correcting the marginal removes this coverage-dependent component of error, while the residual tracks class separability The effective intervention is instead to change the class-level mass allocation itself Matching the marginal to the target composition removes this capacity-driven component and reveals the residual geometric error, but the correction is useful primarily where coverage is limiting and requires an estimate of the target composition The source marginal imposes a class-level transport capacity, so a class can be starved when its target demand
Improvements for AI systems
- Bold header: Correcting Marginal Capacity Failure
Matching the marginal to the target composition substantially improves label transfer
by setting source mass proportional to target demand rather than source frequency, which is crucial when source–target composition shifts can cause rare classes to be systematically under-predicted even with high transport confidence.
- Bold header: Decomposing Transfer Error
The system should be designed to separate errors into a removable component
and a residual component,
where the removable part is governed by coverage and the residual error is associated with class separability,
thus distinguishing between representation quality and capacity failure.
- Bold header: Targeted Composition Estimation Integration
The improved system should incorporate an estimator for target composition, as the paper shows that matching the marginal to the target composition substantially improves label transfer
when using a best deployable estimator
like kNN pre-pass or JCPOT.
Sources
- A Gentle Introduction to Conformal Prediction and Distribution-Free Uncertainty Quantification
- Benchmarking Transcriptomics Foundation Models for Perturbation Analysis : one PCA still rules them all
- Deep Learning Scaling is Predictable, Empirically
- Scaling Laws for Neural Language Models
- Sinkhorn Divergences for Unbalanced Optimal Transport
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks