vislearn/FrEIA

Retrieving Log Jacobian Determinant (ljd) between two specific nodes

mtailanian opened this issue · 4 comments

Hi, I'd like to ask if there is any way I can retrieve the flowing ljd between a specific output and another specific node in the graph.

That would be the short question, but I'll try to give more context, in case it's useful.


I have a nicely defined graph, with one input and multiple outputs.
For simplicity, suppose in some part of the network I perform N invertible downsamplings, and after each one I have a sequential flow.

From the flow I only obtain 1 log jac det, like this:

(z1, z2, ..., zN), ljd = flow(input_volume)

I want to compute the probability just before each mentioned sequential flow. But if I do it like this:

prob_i = 0.5 * torch.sum(z ** 2, dim=(1, 2, 3)) - ljd

it wouldn't be right, because ljd corresponds to the whole flow. Instead I'd need something like (ljd_i, ljd2, ..., ljdN), taken at those specific nodes where I want to evaluate the probability.

Thanks so much in advance, I hope I could explain my needs successfully!

Any help is welcomed!
Thanks again

Hi, I'm not sure that I completely understand your question, but I think you have a misconception about how you can use the log determinants. If you have some invertible layers and then split the input, there is in general no way to split up the Jacobian determinant into two independent parts corresponding to the two outputs. Let $f_1(x), f_2(x)$ be the two outputs of the invertible function. What you have is basically:

$J = \left[ J_1, J_2 \right] = \left[ \frac{\partial f_1(x)}{\partial x}, \frac{\partial f_2(x)}{\partial x} \right]$

If we try to split this up into the corresponding parts, we find that we cannot calculate $\log \det J$ from considering the $J_1$ and $J_2$ terms independently, meaning that we cannot separate the contribution due to the two outputs from each other, except in special cases. Hence your goal to break down the log likelihood into contributions from the different scales in your output is not possible in general. I hope this is clear and helps you with your problem.

Hi @psorrenson,
Thank you very much for taking the time to answer.

I just edited my first comment, trying to make it more clear. Please let me know if it is still confusing...

I understand what you say, you're right, but still, I think it would be possible to retrieve the flowing log jac det in some intermediate node. What do you think? Maybe now with the edited answer my need is more clear...

Thanks again!

Hi @mtailanian, thanks for the clarification. I think I correctly interpreted your question, so my answer still stands: it is not possible to split $\log \det J$ into parts corresponding to your output splits. The only exception would be the following: let $f_1(x), \dots, f_N(x)$ be your outputs and let $J_i(x) = \frac{\partial f_i(x)}{\partial x}$. Then $\log \det J(x) = \sum_i \frac{1}{2} \log \det(J_i(x) J_i(x)^T)$ if $J_i(x) J_j(x)^T = 0$, $i \neq j$. In this case you could decompose $\log \det J(x)$ into the components of the sum. However, this is very unlikely to be the case for all except a few $x$. I encourage you to explicitly write out a formula for your network architecture or a part of it to verify this, or explicitly calculate these quantities using autograd to see that it's the case.

Hi @psorrenson Thanks so much again for the explanation! I think now I understood your answer.
I'll do the math to see if that condition stands, it could be the case!
Thanks again!
I'm closing the issue for now