Great questions!
We explored alternatives to KD-tree search, including alternative tree search methods and LSH. The corresponding code is in this folder in our monty_lab repo. Particularly, the KNNSearch.ipynb contains relevant code for LSH at the bottom. However, it didn’t turn out to be more efficient. It seems like for search in 3D space (as opposed to higher dimensional space) kd-tree search is already quite optimized and LSH used a lot more memory while not providing much speedup. Later on when I put together the constrained object models (GridObjectModel in the code) I again tried alternative search by indexing the grid but again, it wasn’t faster. However, I am not an expert with this so I might have missed something.
The current kd-tree lookup is implemented here: tbp.monty/src/tbp/monty/frameworks/models/object_model.py at 03b8e2211d8acc1e6e30d3330bce3d75605fb8da · thousandbrainsproject/tbp.monty · GitHub (the comment below refers to the alternative I tried). If we work with pretrained models, the KD tree is only built once (tbp.monty/src/tbp/monty/frameworks/models/object_model.py at 03b8e2211d8acc1e6e30d3330bce3d75605fb8da · thousandbrainsproject/tbp.monty · GitHub) and then repeatedly queried at every step (giving it all the search locations for one object at once). If we do continual learning, the tree is rebuilt every time the model in memory is updated (another reason why LSH wasn’t a great fit because the longer-term view of Monty is that it is always learning).
We ran some hyperparameter searches for num_workers, leafsize, and different kd_tree implementations of different libraries.
Definitely happy to hear about any further ideas you try out!
- Viviane