Summary

Resolve #252

Details

JSD

We expect input $X$ and target $Y$ are distributions in log-space, i.e., $X = log Q$ and $Y = log P$.
Jenson-Shannon Divergence between two distributions $P$ and $Q$ is defined as:

where $M = \frac{1}{2}(P + Q)$ is the average distribution and $KL$ is the Kullback-Leibler divergence.
Given that $X = log Q$ and $Y = log P$, we can simplify JSD expression to:

We define the point-wise JSD as:

With point-wise JSD, it's easier to implement JSDs with respect to different reduction methods in future.
The only downside is that it creates a torch.float32 tensor with the same shape as input's.

Current implementation is hardcoded to batchmean which is the original JSD definition.

Gradients

Given:

where $Q = e^X$, $P = e^Y$, and $M = \frac{1}{2}(e^X + e^Y)$.

Gradients of $KL(P\ \Vert\ M)$ with respect to $X_i$:

Gradients of $KL(Q\ \Vert\ M)$ with respect to $X_i$:

Final gradients of JSD:

Combine the results from two KL divergence terms:

Simplify this to:

We store gradients at X_ptr in forward pass to save memory, then retrieve it through ctx in backward function as cross_entropy does. (inplace)

note: inplace operations on inputs might cause an issue with gradient computation.

Testing Done

With inplace (Storing gradients to inputs)

reduce memory usage by 61.54%

increase speed by 53.64%

Without inplace

reduce memory usage by 53%

increase speed by 61%

• Hardware Type: H100
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence