improved multihead pretrained set selection for fine tuning

ACEsuit / mace

MACE - Fast and accurate machine learning interatomic potentials with higher order equivariant message passing.

Other

412 stars 155 forks source link

improved multihead pretrained set selection for fine tuning #451

Open bernstei opened 3 weeks ago

bernstei commented 3 weeks ago

Major refactor of code that subselects pretrained configurations for multihead playback finetuning

Lots of cleanup of code, and substantial change in strategy. For now:

First pick by requested strategy
If it's too many, add only a subset
- first select subset that matches "exclusive" (perfect element match)
- then select subset that matches "combinations" (elements in config are a subset of desired)
- then select subset that matches "any" (any overlap between elements in config and desired)
- if at any point there are too many, pick subset weighted toward better match of elements
If that's not enough, add more with FPS

bernstei commented 2 weeks ago

@ilyes319 thoughts about this one? I think that, even if this isn't exactly the heuristic we use long term, the changes I made are an improvement to the overall design, and will be easier to work with in the future.

ilyes319 commented 2 weeks ago

@bernstei It looks good but I just realized that installing fpsample on CSD3 is a mess, so I am thinking we should find another way to do that part.

bernstei commented 2 weeks ago

So pip install isn't working? Anyway, I'm happy to discuss. I've implemented a naive FPS, but it's not going to be feasible for thousands of points and 2e4 dimensions. I.e., we'll need something more sophisticated. I can imagine dimensional reduction (kernel-PCA?) being part of it, but even with that, picking thousands out of tens of thousands would require some of the complex data structures/heuristics.

Note that we could also do it in two PRs. Accept this one, maybe with small additional design choices to make replacing fpsample easier (it should already support random instead of FPS), or as is, and doing the actual implementation replacement later.