Graph training project for Apple Silicon helped by mlx-graphs

An image generated with Image Creator from Microsoft Designer

One of the problems of creating a machine learning project is training the model on large datasets. This relies on a lot of computing power to chew through the data, but improvements here can help speed up training, and potentially improve the models.

A new project from PHD student Tristan Bilot, Francesco Farina, and the MLX team, mlx-graphs is a library intended to help Graph Neural Networks (GNNs) to run more efficiently on Apple Silicon. GNNs are used to make predictions of nodes, edges, and in performing graph-based tasks, with a particular usefulness in computer vision.

Based on MLX, the mlx-graphs project has been released as a Graph Neural Network library specifically for Apple Silicon. To researchers in the field, the project aims to provide a considerable performance boost.

Bilot claims that initial benchmarks for the library can run at up to ten times the speed of frameworks such as PyTorch Geometric and DGL when training on large graph datasets. It does so by using dedicated kernels designed to parallelize GNN computations running directly on the M-series chip’s GPU.

The project is still in its early days, having only been worked on for a few weeks, with Bilot admitting that there is “still plenty of room for major contributions.” This could be a hint that more speed gains could be found with further development.

The mlx-graphs library is available on GitHub to download and install. Bilot has offered an invitation for others to explore and test the library, provide feedback, and to submit implementations via pull requests.

The project is part of a wave of interest in machine learning and generative AI, a field that can massively transform the creation of content, and the serving of information to users.

Apple CEO Tim Cook has also spoken about the massive effort put into AI features that Apple will be rolling out to users later in 2024.