[<< | Prev | Index | Next | >>] Sunday, July 19, 2026
Implementing PDM
[Research Log: universal math for self-supervised learning--implementation details.]
Pursuant Probability Density Magic, the following implementation notes aren't conceptually new since then but some might find these useful. Recall equation (1):
` del_omega log p_omega(x) = del_omega log p_z(z) + (: del_omega dz/dx , {: dx/dz :}^T :)_F` (1) Where `(: ,:)_F` is the Frobenius inner product* and (1) tells us how to maximize the probability of our model generating a training set. I.e., it is an ideal, unsupervised learning algorithm.
This is mostly linearly separable by layer, except for where non-linear layers are location sensitive. In fact, once we fully account for that location sensitivity, the equation gets, in some sense, simpler:
` del_omega log p_omega(x) = (: del_omega dz/dx , {: dx/dz :}^T :)_F` (2) In order to do this, we simply require that `z` is sampled uniformly from any finite region (which makes `del_omega log p_z(z) = 0`), which we can do by replacing our final PDF with a non-linear layer that maps through the integral of that PDF. E.g., if our original PDF on `z` is normal, then we can instead map our output layer through `Phi`, which is just a sigmoidal function that integrates to exactly the normal distribution, and now `z` is uniform in the range 0 to 1 and we can drop the first term in (1).
But this begs the question of how to actually handle that non-linearity/location sensitivity.
For that, to be done efficiently, we essentially replicate reverse mode differentiation (back propagation) on the location term, which with some futzing can be done with existing auto gradient methods but I will spell it out here because it's useful to understand intuitively and also suggests some easy direct implementations.
In the general form, for some single middle layer mapping `x` to `z`, (2) becomes:
` del_omega log p_omega(x) = (: del_omega dz/dx , {: dx/dz :}^T :)_F + S_"in" \cdot dz/(d omega)` (3) This is for any `omega` which is a parameter of that layer. By passing `S` backward down the chain, we are able to fully modularize our gradient calculations this way--the gradient against a given `omega` is entirely calculated within the layer that holds it.
Now we chain S backwards like this:
` S_"out" = S_"in" dz/dx + (: del_x dz/dx , {: dx/dz :}^T :)_F` (4) Note `del_x` here, not `del_omega`.
As a vector, `S` is the size of `z` and `x` and simply represents the gradient at `z` of `log p`. It is, in effect, the implicit `d/dz log p_z(z)` component of the `del_omega log p_z(z)` we dropped from (1), but now propagated backward to represent all remaining layers instead of just the final output one.
Now I'll simplify for some common cases:
For any linear (not location sensitive) layer, `S_"out"` simplifies to:
` S_"out" = S_"in" dz/dx` (5) So, for instance, a matrix multiply layer gives:
` del_M log p_omega(x) = M^-T + S_"in" x^T` (6)
` S_"out" = S_"in" M` (7)
For any non-parameterized, element-wise non-linearity `z=phi(x)`, we get (assuming element-wise operators):
` S_"out" = S_"in" phi'(x) + (phi''(x))/(phi'(x))` (8) Now it's (marginally) worth noting that if `phi(x)` is a CDF of some PDF `pi(x)`, such that `phi'(x) = pi(x)`, then:
` S_"out" = S_"in" pi(x) + d/dx log pi(x)` (9) That is, if our squashing function corresponds to a PDF with a known log derivative, our work is done:
For `pi` = a normal distribution:
` S_"out" = S_"in" pi(x) - x` (10) For `pi` = Laplace:
` S_"out" = S_"in" pi(x) - "sign"(x)` (11) For `phi` = logistic sigmoid (`pi` = logistic distribution):
` S_"out" = S_"in" pi(x) + (1 - 2z)` (12) However, I say "marginally" worth noting here because bounded squashing functions (any CDF) make poor candidates for inversion...
Finally, for an arbitrary final output PDF (the term we dropped from (1)), we can just introduce it as an originating `S`:
` S_o = grad_z log p_z(z)` (13) [<< | Prev | Index | Next | >>]