Is your feature request related to a problem? Please describe.
Currently, it is difficult or impossible to properly regularize latent spaces because there is no loss function available to compare and align distributions in the latent domain. While reconstruction and task-specific losses can constrain individual samples, they do not provide a mechanism to enforce desirable global properties of the latent distribution.
As a result, latent representations may become poorly structured, drift between training iterations, or fail to match target distributions. This limitation makes it challenging to implement distribution-based regularization strategies commonly used in generative modeling, domain adaptation, and representation learning.
Describe the solution you'd like
Implement a Sinkhorn loss based on entropic optimal transport to measure discrepancies between latent distributions.
The Sinkhorn loss would allow users to:
Regularize latent spaces by aligning learned distributions with target distributions.
Enforce consistency between latent representations from different datasets, domains, or models.
Improve stability and structure of learned embeddings.
Leverage optimal transport distances while maintaining computational efficiency through Sinkhorn iterations.
Ideally, the implementation would support batched computation, GPU acceleration, and configurable regularization parameters (e.g., entropy regularization strength and number of Sinkhorn iterations).
Describe alternatives you've considered
Several alternatives can partially address this issue, but each has limitations:
KL divergence requires explicit density estimation or parametric assumptions about the latent distributions.
Maximum Mean Discrepancy (MMD) can compare distributions but may be sensitive to kernel selection and may not capture geometric structure as effectively.
Adversarial losses can align distributions but introduce additional model complexity and training instability.
Manual regularization terms based on latent statistics (mean, variance, covariance) only constrain low-order moments and do not fully compare distributions.
Compared to these approaches, Sinkhorn loss provides a principled and flexible optimal-transport-based metric that captures the geometry of latent distributions while remaining computationally tractable.
Additional context
Sinkhorn loss has become a standard tool in representation learning and generative modeling for comparing empirical distributions. Its availability would enable more advanced latent-space regularization strategies and broaden the range of models that can be implemented within the framework.
Is your feature request related to a problem? Please describe.
Currently, it is difficult or impossible to properly regularize latent spaces because there is no loss function available to compare and align distributions in the latent domain. While reconstruction and task-specific losses can constrain individual samples, they do not provide a mechanism to enforce desirable global properties of the latent distribution.
As a result, latent representations may become poorly structured, drift between training iterations, or fail to match target distributions. This limitation makes it challenging to implement distribution-based regularization strategies commonly used in generative modeling, domain adaptation, and representation learning.
Describe the solution you'd like
Implement a Sinkhorn loss based on entropic optimal transport to measure discrepancies between latent distributions.
The Sinkhorn loss would allow users to:
Regularize latent spaces by aligning learned distributions with target distributions.
Enforce consistency between latent representations from different datasets, domains, or models.
Improve stability and structure of learned embeddings.
Leverage optimal transport distances while maintaining computational efficiency through Sinkhorn iterations.
Ideally, the implementation would support batched computation, GPU acceleration, and configurable regularization parameters (e.g., entropy regularization strength and number of Sinkhorn iterations).
Describe alternatives you've considered
Several alternatives can partially address this issue, but each has limitations:
KL divergence requires explicit density estimation or parametric assumptions about the latent distributions.
Maximum Mean Discrepancy (MMD) can compare distributions but may be sensitive to kernel selection and may not capture geometric structure as effectively.
Adversarial losses can align distributions but introduce additional model complexity and training instability.
Manual regularization terms based on latent statistics (mean, variance, covariance) only constrain low-order moments and do not fully compare distributions.
Compared to these approaches, Sinkhorn loss provides a principled and flexible optimal-transport-based metric that captures the geometry of latent distributions while remaining computationally tractable.
Additional context
Sinkhorn loss has become a standard tool in representation learning and generative modeling for comparing empirical distributions. Its availability would enable more advanced latent-space regularization strategies and broaden the range of models that can be implemented within the framework.