Repository: pytorch
Issue: Stabilizing the Kumaraswamy Distribution
Labels: module: numerical-stability,module: distributions,triaged

### 🐛 Describe the bug

The current implementation of the Kumaraswamy distribution has significant numerical instabilities in its log-pdf, inverse CDF, and their parameter gradient expressions. All such expressions contain unstable terms of the form `log(1 - exp(x))`. I believe these instabilities have significantly hindered the usage of the Kumarswamy.

I recently published a paper [Stabilizing the Kumaraswamy Distribution](https://arxiv.org/abs/2410.00660) which diagnoses these issues and provides a stable implementation with an unconstrained logarithmic parameterization.

My implementation inherits from `Distribution`, not `TransformedDistribution`; relying on the sequence of provided transformations applied to the base distribution, namely 

```
transforms = [
            PowerTransform(exponent=self.concentration0.reciprocal()),
            AffineTransform(loc=1.0, scale=-1.0),
            PowerTransform(exponent=self.concentration1.reciprocal()),
        ]

```
produces the aforementioned `log(1-exp(x))` instabilities.


If there is interest in correcting the PyTorch implementation, I am happy to help, I just ask the paper be cited in the documentation.

Let me know how we should proceed.


-------

Some Evidence of the communities struggle with these instabilities
- [torch.distributions.kumaraswamy.Kumaraswamy generates samples outside its support (0,1)](https://github.com/pytorch/pytorch/issues/95548)
- [Implement Kumaraswamy Distribution](https://github.com/probtorch/pytorch/pull/143)

See also Section 5: Related Work, 'Kumaraswamy as a Beta surrogate' in my paper. It briefly goes over the hacks required in prior works using the Kumaraswamy (across Pytorch, Tensorflow, Theano, etc), typically relying on clipping of KS parameters (thus limiting expressiveness of the distribution class) and/or constraining the standard uniform base distribution (and thus distorting the sampling distribution).


cc @fritzo @neerajprad @alicanb @nikitaved

### Versions

python collect_env.py
--2024-10-27 18:06:10--  https://raw.githubusercontent.com/pytorch/pytorch/main/torch/utils/collect_env.py
Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.110.133, 185.199.111.133, 185.199.108.133, ...
Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.110.133|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 24366 (24K) [text/plain]
Saving to: ‘collect_env.py’

collect_env.py                   100%[========================================================>]  23.79K  --.-KB/s    in 0.03s   

2024-10-27 18:06:10 (810 KB/s) - ‘collect_env.py’ saved [24366/24366]

Collecting environment information...
PyTorch version: 2.4.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.0.1 (arm64)
GCC version: Could not collect
Clang version: 14.0.0 (clang-1400.0.29.202)
CMake version: Could not collect
Libc version: N/A

Python version: 3.9.19 | packaged by conda-forge | (main, Mar 20 2024, 12:55:20)  [Clang 16.0.6 ] (64-bit runtime)
Python platform: macOS-13.0.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M2

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] optree==0.11.0
[pip3] pytorch-lightning==2.2.2
[pip3] torch==2.4.0
[pip3] torch_geometric==2.5.3
[pip3] torchmetrics==1.4.0
[pip3] torchrl==0.5.0
[pip3] torchvision==0.16.1
[conda] libtorch                  2.1.2           cpu_generic_h811e168_4    conda-forge
[conda] nomkl                     1.0                  h5ca1d4c_0    conda-forge
[conda] numpy                     1.26.4                   pypi_0    pypi
[conda] optree                    0.11.0                   pypi_0    pypi
[conda] pytorch-lightning         2.2.2              pyhd8ed1ab_0    conda-forge
[conda] torch                     2.4.0                    pypi_0    pypi
[conda] torch-geometric           2.5.3                    pypi_0    pypi
[conda] torchmetrics              1.4.0              pyhd8ed1ab_0    conda-forge
[conda] torchrl                   0.5.0                    pypi_0    pypi
[conda] torchvision               0.16.1          cpu_py39h936d113_3    conda-forge