## Summary Resolve linkedin#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: ```math JSD(X, Y) = JSD(P\ \Vert \ Q) = \frac{1}{2} (KL(P\ \Vert\ M) + KL(Q\ \Vert\ M)) ``` 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: ```math \begin{align} JSD(X, Y) &= \frac{1}{2}(\sum_i P_i\ log\frac{P_i}{M_i} + \sum_i Q_i\ log\frac{Q_i}{M_i})\ &= \frac{1}{2} \sum_i (P_i\ log\ P_i - P_i\ log\ M_i + Q_i\ log\ Q_i - Q_i\ log\ M_i)\ &= \frac{1}{2} \sum_i (P_i\ log\ P_i + Q_i\ log\ Q_i - 2M_i\ log\ M_i)\ &= \frac{1}{2} \sum_i (P_i \cdot X_i + Q_i\cdot Y_i - 2M_i\ log\ M_i) \end{align} ``` We define the point-wise JSD as: ```math JSD(X_i, Y_i)= \frac{1}{2} (P_i \cdot X_i + Q_i\cdot Y_i - 2M_i\ log\ M_i) ``` 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: ```math JSD(X, Y) = JSD(P\ \Vert \ Q) = \frac{1}{2} (KL(P\ \Vert\ M) + KL(Q\ \Vert\ M)) ``` 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$: ```math \begin{align} \frac{\partial}{\partial X_i} \sum_j P_j\ log\frac{P_j}{M_j} &= \frac{\partial}{\partial X_i} \sum_j P_j (Y_j - log\ M_j)\ &= \frac{\partial}{\partial X_i} \sum_j - P_j log\ M_j\ &= - P_i \cdot \frac{1}{M_i}\cdot \frac{e^X}{2} = -P_i\cdot \frac{Q_i}{2M_i} \end{align} ``` #### Gradients of $KL(Q\ \Vert\ M)$ with respect to $X_i$: ```math \begin{align} \frac{\partial}{\partial X_i} \sum_j Q_j\ log\frac{Q_j}{M_j} &= \frac{\partial}{\partial X_i} \sum_j Q_j (X_j - log\ M_j)\ &= \sum_j\left( \frac{\partial Q_j}{\partial X_i}(X_j - log\ M_j) + Q_j \frac{\partial (X_j - log\ M_j)}{\partial X_i}\right)\ &= Q_i(X_i - log\ M_i) + Q_i (1 - \frac{Q_i}{2M_i}) \end{align} ``` #### Final gradients of JSD: Combine the results from two KL divergence terms: ```math \begin{align} \frac{\partial}{\partial X_i} JSD(X, Y) &= \frac{1}{2}\left(-P_i\cdot \frac{Q_i}{2M_i} + Q_i(X_i - log\ M_i) + Q_i (1 - \frac{Q_i}{2M_i})\right) \end{align} ``` Simplify this to: ```math \begin{align} \frac{\partial}{\partial X_i} JSD(X, Y) &= \frac{1}{2}\left(-P_i\cdot \frac{Q_i}{2M_i} + Q_i(X_i - log\ M_i) + Q_i (1 - \frac{Q_i}{2M_i})\right)\ &= \frac{1}{2}\left(Q_i (X_i - log\ M_i + 1 - \frac{P_i + Q_i}{2M_i})\right),\ where\ 2M_i = P_i + Q_i \ &= \frac{1}{2}\cdot Q_i \cdot (X_i - log\ M_i) \end{align} ``` 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% ![jsd_memory](https://github.com/user-attachments/assets/fc71a3b3-73aa-433f-ac62-0a0924d5c2de) increase speed by 53.64% ![jsd_speed](https://github.com/user-attachments/assets/98eac714-aee3-4326-8554-26efc473baac) ### Without inplace reduce memory usage by 53% ![jsd_memory](https://github.com/user-attachments/assets/4f435fdb-872e-4465-a086-9ebcaa41ba3f) increase speed by 61% ![jsd_speed](https://github.com/user-attachments/assets/4e88bccd-982e-4fca-9148-004e4ef97570) - Hardware Type: H100 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence --------- Co-authored-by: Shao Tang <tangshao28@gmail.com> Co-authored-by: Qingquan Song <ustcsqq@gmail.com>

reverse KL and JSD #252

Description

🚀 The feature, motivation and pitch

Alternatives

Additional context

Metadata

Metadata

Assignees

Labels

Type

Fields

Projects

Milestone

Relationships

Development

Issue actions