eepperly / XTrace

Code for the paper "XTrace: Making the most of every sample in stochastic trace estimation"
MIT License
10 stars 1 forks source link

Sanity Check #2

Open quattro opened 3 months ago

quattro commented 3 months ago

Hi!

Really cool work on stochastic trace estimation and implementation.

I have a basic question regarding the 'improved' implementations. Namely, to rescale the sampled test-vectors both XTrace anx XNysTrace perform something along the lines of,

...
S = cnormc(inv(R)');
scale = (n - m + 1) ./ (n - (vecnorm(W)').^2  + abs(diag_prod(S,W) .* vecnorm(S)').^2);

Is the point-wise multiplication between diag_prod(S, W) with vecnorm(S) necessary? S is computed by standardizing inv(R)', so vecnorm(S)' should be 1. Is that correct, and if so, I suppose the additional multiplication can be dropped?

quattro commented 3 months ago

PS: I've been toying around with some implementations in Python using JAX and Lineax here, if you're curious: https://mancusolab.github.io/traceax/

Cheers and thanks again.