Abnormal Memory Demand

[ben0i0d]

When using t-SNE on a dataset of shape (2500, 784) on the CPU, I found that the memory usage peaked at around 34.5GB, which made it impossible to run the program on my local computer.

This is quite unreasonable since the dataset is small. Furthermore, I did not observe such behavior with sklearn or pcax; they showed almost no memory demand.

I have uploaded my code and process to https://github.com/ben0i0d/tsne-test.

If you have time, please take a look.

[ben0i0d]

I am using

JAX version 0.4.37
pcax version 0.1.0
tsnex version 0.0.1

installed via pip.

[ben0i0d]

related to

X_new, _ = jax.lax.scan(train_step_early_exageration, X_new, xs=None, length=n_exageration)  
X_new, _ = jax.lax.scan(train_step, X_new, xs=None, length=n_iter)  

[alonfnt]

Hi @ben0i0d
Thanks for raising this and thanks a lot for the repo to test it.

I agree with you that it is unresonable, but sadly it is working "as expected".
The current implementation (which is quite simple since I needed it to compute some gradients), relies on computing the distance matrix directly. You can get the same error if you run

dist_fn = jax.vmap(jax.vmap(tsnex.tsne.euclidean_distance, in_axes=(0, None)), in_axes=(None, 0))
D = dist_fn(x,x)

with some napkin calculations: assuming single floating point (32bits), to compute D we need to create a (2500, 2500, 784) matrix, which should give ~20Gb, then we add D and we get the >30Gb.

The reason sklearn can get away with it is because, other than all the approximations method that I need to include #4 #5 , they rely on scipy distance calculations (when using method='exact' as an argument, which would be equivalnet to tsnex), and scipy has a neat optimization for square matrices distances.

Now, what can we do... The first simple solution could be to check the dimension of x, and if above a threshold k, split D to compute every k rows (that would save some memory).

Ideally, I think more optimize distance calculations should be implemented, but that takes time.

If you feel to contribute I'd gladly review the above "quick" solution, otherwise I'll give it a go soon.

Thanks for pointing this out and let's try to solve it!

[ben0i0d]

Thank you for your reply. I also noticed related to the distance computation:
metric_fn = jax.vmap(jax.vmap(metric_fn, in_axes=(0, None)), in_axes=(None, 0))

If I have time later, I would be interested in learning more about the optimization methods you mentioned （#4 #5） and trying some code.

Since I am working on an ROCm platform and almost all tSNE acceleration implementations are written for CUDA, I also tried the tSNE-Torch implementation. Unfortunately, it took 2 minutes and 13 seconds on this dataset, whereas this JAX implementation is extremely fast, taking only 0.1 seconds (so I truly appreciate your work).