Using mlverse/torch on spatial omics, e.g. for applying foundation models

[Artur-man]

CC @JuanHenaoS

Mapping pretrained weights of sc/spatial foundation models to torch:
https://github.com/mlverse/torch
https://torch.mlverse.org/docs/articles/python-to-r

[jdhenaos]

Hi!

Thanks for the links!

That improves the workload a lot!

Could you explain the idea of benchmarking in more detail?

Implementing foundational models in R and comparing them to the previously established ones? runtime? vs reticulate? (I am guessing)

[Artur-man]

No actually, the idea would be to load the pretrained weights of a (ideally sc/spatial) foundation model in the R wrapper (mlverse/torch) and only to generate the embeddings of the test data.

The benchmarking would be to compare the same embedding generation in Python and R, and show/prove/test the runtime and conformance (whether you can generate the same embedding). If there are bottlenecks, report ... open issues in mlverse/torch so the developers (together with us) work on improving.

The motivation is to show the community that you can use the model without using python frameworks, and potentially make the models part of BioC workflows.

We can compare with reticulate as well.

[jdhenaos]

awesome!

I think we can start with one foundational model during the hackathon, including defining the benchmarking datasets, experimental setup, and metrics.

What about nicheformer?: https://github.com/theislab/nicheformer

[Artur-man]

one thing to consider is that we may need additional support for transformers layers, which i am not %100 sure if they are available in the main torch(mlverse) implementation ... and even if it is, I do not know if we can map pre-trained weights of an architecture with transformers layer.

We can also aim for simple NN models to start with. But again, I am a fan of discussions and experimentation, we do not need a crystal clear plan before :D ... that being said, lets keep discussing here if we have time.

Note: I will be in another hackathon until Venice.

[jdhenaos]

Totally agree.

I was scaling things too fast.

It is necessary to proceed from the simplest possible scenario to first get familiar with mlverse/torch and have plenty of control over the code and expected results.

How do you think we should proceed for the very first steps? Looking for basic NNs in single-cell? Train our own MLPs in Python, then compare the embeddings to R? Or could it be setting up a Dockerfile to provide a ready-to-go container before the hackathon?

You can reply at your convenience during the hackathon. I am more flexible with the time, so I can reply quickly

[Artur-man]

Looking for basic NNs in single-cell?

Yes I think this is a really good idea ... we need a simple architecture which we can rewrite in mlverse/torch, and yes depending on the model we

1. either parse pre-trained weights from the repo where the model is available, or
2. train it ourselves, generate weights ... load to mlverse/torch compare.

I am not sure if we need a docker ? ... what OS and laptop you are using @jdhenaos? I have a macOS and MPS gpu device, so I might run GPU locally (we may not need GPU but who knows). If you do not I can try arranging a remote one.

Ok lets continue discussing here, i would suggest experimenting with the package until we go to Venice. I will try doing the same, I have a student who has a code base for this. I will ask him to give me that too.

[jdhenaos]

I have a Linux with an NVIDIA RTX 500 Ada GPU. So we can both run locally (if necessary).

I will play around with the package. Having a code base will be very useful as well.

[jdhenaos]

Happy to see this idea is catching attention!

I think the discussion outlines the to-dos for the hackathon and the next steps overall to demonstrate the project's relevance to the community.

I think we can run simple NNs to test mlverse/torch's capabilities, along with testing improvements that Rarr could carry. Also, test Rarr against mlverse/torch for some simpler tasks (I guess the logic is obj -> zarr -> model, not zarr inside the model by itself (??)).

We can create a minimal example using sc/spatial data with a simple NN and demonstrate the usability of mlverse/torch when integrated with R development, and compare the process when using Python interoperability with the same aim (which I think is one of the major contributions).

[Artur-man]

By the way i mimicked the benchmarking in https://github.com/mlverse/torch/blob/main/bench/small-nn.R

See: https://github.com/Artur-man/py_r_torch_benchmarking

(py_r_torch_benchmarking) LANDTHALER57:py_r_torch_benchmarking amanuky$ python benchmark_device_simple/benchmark_device_simple.py 
{'load': 0.0013628005981445312, 'setup': 0.0018796920776367188, 'mps': 1.6316609382629395, 'cpu': 6.573911190032959}
{'load': 0.0011260509490966797, 'setup': 0.0017371177673339844, 'mps': 1.5132861137390137, 'cpu': 6.622015953063965}
{'load': 0.0011050701141357422, 'setup': 0.002045869827270508, 'mps': 1.5158348083496094, 'cpu': 6.62052583694458}
{'load': 0.0010769367218017578, 'setup': 0.0025539398193359375, 'mps': 1.518923044204712, 'cpu': 6.630447149276733}
{'load': 0.001125335693359375, 'setup': 0.002821207046508789, 'mps': 1.5086002349853516, 'cpu': 6.6927337646484375}
(py_r_torch_benchmarking) LANDTHALER57:py_r_torch_benchmarking amanuky$ Rscript benchmark_device_simple/benchmark_device_simple.R
  mps   cpu  load setup 
 1.77  9.83  0.00  0.07 
  mps   cpu  load setup 
 1.68  9.74  0.00  0.08 
  mps   cpu  load setup 
 1.78  9.73  0.00  0.07 
  mps   cpu  load setup 
 1.76  9.61  0.00  0.07 
  mps   cpu  load setup 
 1.77  9.80  0.00  0.07 

It seems there is native mps/cuda support in torch without python thats nice .... and it accelerates the model as much as the python version does. It is a bit slow though but not surprising.

[jdhenaos]

You can work natively with mps/cuda in mlverse/torch:

https://torch.mlverse.org/docs/articles/installation

For mps, it is not a big deal (apparently), but for CUDA, I will fix in advance the version compatibility. mlvererse/torch works with v12, and I have v13.

I am not surprised that R takes a bit more time, but providing GPU usability in R will be widely welcomed for the community, overall for Nicholas Cooley and the rest of the people interested in opening a dedicated GPU bioc group.

Maybe we can show the advances after the Hackathon. I think that will be valuable for Nicholas, who is organizing a GPU-dedicated group for the EuroBioC hackathon in June.

[Artur-man]

Amazing! is nicholas joining venice too or ?

[jdhenaos]

He won't be there afaik.

But I will attend the Hackathon in June (virtually). Will you? Otherwise, I can keep in touch with you during those days to share the news and conclusions.

In advance, we can share the results from next week's hackathon with him.

[gbrsales]

Hi! This looks like a great start, and the initial benchmarking results are very encouraging.

One thing I wanted to note is that while mapping the pretrained weights addresses the structural gap, another significant challenge will likely be the data pipeline.

For sc/spatial foundation models (like Geneformer or scGPT), the preprocessing is often as complex as the model itself. If the R implementation of things like tokenization, rank-value encoding, and dynamic masking doesn't mirror the Python preprocessing logic exactly, the pretrained weights won't produce the same results as the original code.

While tools like PyTorch tracing can automate the 'export' of the neural architecture, they don't help with the data preparation. Re-implementing the custom torch::dataset and torch::dataloader logic in R would still be a manual process and likely the primary bottleneck.

Has anyone here had experience porting complex Python data-loading pipelines to mlverse/torch? I'd love to hear how you handled the tokenization/encoding side of things.