Decision trees are a commonly used machine learning technique. They often are preferred to neural networks or more complex models for small problems that wouldn’t need deep learning; libraries implementing decision tree “forests” trained via gradient boosting like XGBoost and CatBoost are some of the most popular ML frameworks out there.

A decision tree generally works as follows: given your input feature vector $\v{x}$, each binary node $i$ computes a linear function $\v{v}_i\cdot\v{x}+b_i$. If this comes out > 0, go one side, otherwise go the other. Once you’re at a leaf, you can read out its related outcome, and that’s your answer. There are variations; forests for example aggregate multiple trees to obtain a more robust judgement. But the core principle stays the same.

A few months ago I released the first version of MiniML, a small machine learning framework powered by Jax. The main idea behind it was to have a very lean toolset to build models that would be as simple to train as a Scikit-learn one, while offering the same level of flexibility as PyTorch. I wrote all about it here.

A few versions later I’ve expanded on that base with a few much needed basic modules for machine learning - radial basis function networks, multi-head self attention, a testing system that uses gold standard data from Backblaze, support for non-Scipy optimizers and more. In this release though I want to focus on the addition of two features that support research in one particular phenomenon: “grokking”. Let’s see what is it!