Fastest way to compute Empirical Fisher Information Matrix Vector product?

[YouJiacheng]

My draft:

import jax
import jax.numpy as jnp

def log_p(θ, x): # simple example
    return θ @ x

def fvp(θ, xs, dθ):
    @jax.vmap
    def sample_fvp(x):
        _, f_jvp = jax.linearize(lambda θ: log_p(θ, x), θ)
        result_jvp = f_jvp(dθ)
        f_vjp = jax.linear_transpose(f_jvp, dθ)
        result_vjp = f_vjp(result_jvp)[0]
        return result_vjp
    
    return jax.tree_map(lambda x: jnp.mean(x, 0), sample_fvp(xs))

Intention: Using conjugate gradient descent to solve FΔθ=dL/dθ for natural gradient descent.
Fis Fisher Information Matrix, i.e. Expectation(over x) of self outer product of ∇_θ(log_p(x)).

[joeryjoery]

Hey, what you do is I believe quite efficient. I do something similar. However, instead of doing vmap on sample_fvp, you could also vmap over the log_p such that the gradients/ jvp of each sample is independent of one another and that the jacobian is defined as J \in R^(n x d).

Hence the pullback vjp will be cast over the n samples meaning that you can drop the jnp.mean (which should also be a sum if I'm correct) as this pullback aggregates the sample fvp's for you!

I implement it as this (I don't use linearize because I want compatibility with possible aux data).

from typing import Callable, Any, Tuple
from chex import Array, ArrayTree
from jax import vmap, jvp, vjp


def efvp(
        negative_log_likelihood: Callable[[Array, Array], float],  # (outputs, targets) -> (loss)
        model_fun: Callable[[ArrayTree, Array], Array],            # (params, inputs) -> (outputs)
        primals: Any,
        tangents: Any,
        data: Tuple[ArrayTree, ArrayTree],  # D = (x, y) pairs.
        has_aux: bool = False
) -> Tuple[Any, ...]:
    # Ensure that batch-dimensions are independent of one-another during autodiff.
    negative_log_likelihood = vmap(negative_log_likelihood, in_axes=(0, 0))

    # Compose the model and loss function; the antiderivative of the score function.
    x, y = data
    param_loss = lambda p: negative_log_likelihood(model_fun(p, x), y)

    return _core_efvp(param_loss, primals, tangents, has_aux)


def _core_efvp(
        param_loss: Callable,
        primals: Any,
        tangents: Any,
        has_aux: bool = False,
) -> Tuple[Any, ...]:
    y, Jt, *aux = jvp(param_loss, primals, tangents, has_aux=has_aux)
    pullback = vjp(param_loss, *primals, has_aux=has_aux)[1]
    
    if has_aux:
        return y, pullback(Jt), aux
    return y, pullback(Jt)

By the way, if you want to do Natural Gradient descent, the Empirical Fisher may not be the best choice to precondition on. If your model parameterizes an exponential family distribution, you could use the GGN as an alternative, or simply approximate the Fisher conditional on x by drawing samples from model_fun(x) and treating the samples as the y's.

[YouJiacheng]

Thanks! You are right, for jvp and vjp instead of grad or jacobian, vmap before autodiff won't hurt efficiency. But I still think deferring vmap to directly translate mathematical formula into code is better.
(Though it seems convenient to sum the sample fvp's by vjp∘jvp on batched functions, it won't be straightforward to understand.)
(mean v.s. sum is just difference in terminology across fields, I think.)
(BTW, the transpose of a broadcast will be a sum, thus there won't be an actual difference.)

Thanks very much for advice on natural gradient descent: I am a beginner in this field, and effective precondition can be crucial to the success of my work. However, I am working on something like ferminet. I think it isn't an exponential family distribution. Could you provide some references (links or keywords are enough) about "GGN" and "simply approximate the Fisher conditional on x by drawing samples from model_fun(x) and treating the samples as the y's"?

[YouJiacheng]

Okay I find that GGN is generalized Gauss-Newton, I will investigate it!

[joeryjoery]

I think this paper addresses all important points: https://arxiv.org/abs/1905.12558.

I'm not familiar with ferminet, but from the formula you provide I would suspect the Empirical Fisher in this case to correctly approach the Fisher given enough samples of p(x). Though, I don't know what this particular formula would be hiding underneath (I assume p(X) to just be a marginal distribution).

To elaborate:

In the case of an exponential family for your log density (https://en.wikipedia.org/wiki/Exponential_family), it just so happens that the Fisher matrix we care about is equivalent to the gauss-newton matrix under the composition f_loss o f_model. In this special case, the Fisher is equal to: E_p(y,x)[Jac(model)^T @ Hess(loss) @ Jac(model)] where p(y, x) can be approximated by your dataset (average gauss-newton for the full marginal; or a sum of gauss-newtons when conditioned on x).

For anything else, you need to compute an expectation with respect to the model distribution. In the case of a predictive model p_\theta(y | x) we need to marginalize both y and x. Usually it is fine to simply condition on x (just use the data) as we don't have the true distribution of x. But we can marginalize y as the distribution is defined by the model, this is what I meant with sampling from the model. This would go something along the lines of

1. y_hat ~ sample(model(params, x), n=100)
2. augmented_model = lambda p: loss(model(params, jnp.repeat(x, 100), y_hat)
3. Compute fisher for params with augmented model and divide the resulting matrix by the sample size.
4. Vmap 1 to 3 and sum this procedure over each x sample in your dataset.
5. Use this for Natural Gradient Descent :)

[YouJiacheng]

Wow, in fact I just found GGN in this paper!
In ferminet, p(X) is the "full" distribution(not marginal or conditional), there isn't any condition or hidden variable.
Actually it is something like a reinforcement learning, i.e. we need to optimize the distribution p(X) and we can sample from it.
(Of course it will be nice to consider p(X) is conditioned on some Z. And just like what you said, we can simply keep Z in condition, not marginalizing it. Thus it is equivalent to absorb Z into model)

1. y_hat = model(params, x).sample(sample_size)
2. augmented_model = lambda p: model(p, x).evaluate_log_ps(y_hat)
   (BTW, is it ok to simply sum fisher over each x (Z) in dataset? Or we should compute Natural Gradient Descent updates for each x and average these updates?)

[joeryjoery]

Okay, so I would then indeed expect your implementation to be correct :).

In this case, you would average over the x's as you treat them as samples of p(X) over which you are computing the expectation.