MLX: A new era in ML research on Apple Silicon

In the evolving landscape of machine learning research, efficiency and flexibility are important. Meet with MLX, an array framework crafted specifically for harnessing the power of Apple silicon. Developed by Apple’s machine learning research team, MLX is a beacon for researchers keen on pushing the boundaries of machine learning technologies.

What Sets MLX Apart?

MLX prides itself on offering an API that strikes a perfect balance between familiarity and innovation. For those accustomed to NumPy, MLX’s Python API provides a smooth transition, extending support through C++, C, and Swift APIs that reflect its Python counterpart’s design principles. High-level packages such as mlx.nn and mlx.optimizers further simplify complex model construction, mirroring the ease found in PyTorch.

Breakthrough function transformations

Function transformations in MLX are not just an additive feature—they redefine efficiency. With support for automatic differentiation, automatic vectorization, and computation graph optimization, MLX enhances the machine learning research workflow.

Lazy computation and dynamic graph construction

With a focus on optimization, MLX implements lazy computation, only materializing arrays as needed. This approach, combined with dynamic graph construction, alleviates the oft-encountered slow compilation triggered by changes in function argument shapes, making debugging both straightforward and intuitive.

Seamless multi-device support

Gone are the days of cumbersome data transfers between CPU and GPU. MLX’s unified memory model ensures arrays reside in shared memory, allowing operations across supported devices without the extra data transfer overhead.

Designed for and by researchers

The design philosophy behind MLX is clear—simplify research without sacrificing efficiency. The framework’s straightforward structure not only makes it accessible for newcomers but also offers ample opportunities for veterans to innovate and extend its capabilities.

Practical applications and examples

MLX is not a theoretical exercise—it’s a toolkit primed for real-world applications. From training transformer language models to speech recognition with OpenAI’s Whisper, MLX broadens the horizons for machine learning tasks. Its capabilities are further showcased in generating images with Stable Diffusion and fine-tuning large-scale text generation models.

Getting started

Interested in exploring MLX? The framework is readily available on PyPI and conda-forge, ensuring a straightforward setup for your Apple silicon-powered device. Whether you’re working on a Mac with an M series chip or pursuing development with the C++ API, MLX is designed to integrate seamlessly into your workflow.

Implementing Linear Regression with MLX

To demonstrate MLX’s capabilities and ease of use, let’s walk through a straightforward example: implementing a Linear Regression model. This example will highlight MLX’s familiar API, lazy computation, and the framework’s efficiency for machine learning tasks.

Before diving into the code, ensure you have installed MLX on your Apple silicon device. If you haven’t already, follow the installation instructions provided earlier in this article. Once installed, you’re ready to begin.

Step 1: Import MLX and Initialize Variables

Begin by importing the mlx.core module and setting up the necessary metadata for our linear regression problem.

import mlx.core as mx

num_features = 100
num_examples = 1_000
num_iters = 10_000  # Number of iterations for Stochastic Gradient Descent (SGD)
lr = 0.01  # Learning rate for SGD

Step 2: Generate a Synthetic Dataset

Next, we’ll generate a synthetic dataset comprising a design matrix X, a ground truth parameter vector w_star, and dependent values y with added Gaussian noise.

# True parameters
w_star = mx.random.normal((num_features,))

# Input examples (design matrix)
X = mx.random.normal((num_examples, num_features))

# Noisy labels
eps = 1e-2 * mx.random.normal((num_examples,))
y = X @ w_star + eps

Step 3: Define Loss and Gradient Functions

Define the squared loss function and obtain the gradient function of the loss with respect to the parameters.

def loss_fn(w):
    return 0.5 * mx.mean(mx.square(X @ w - y))

grad_fn = mx.grad(loss_fn)

Step 4: Perform Stochastic Gradient Descent (SGD)

Initialize the parameters randomly and repeatedly update them using the gradient descent algorithm for a defined number of iterations.

w = 1e-2 * mx.random.normal((num_features,))

for _ in range(num_iters):
    grad = grad_fn(w)
    w = w - lr * grad
    mx.eval(w)

Step 5: Evaluate the Model

Finally, assess the accuracy of the learned parameters by computing the loss and comparing the learned parameters to the ground truth.

loss = loss_fn(w)
error_norm = mx.sum(mx.square(w - w_star)).item() ** 0.5

print(f"Loss: {loss.item():.5f}, |w-w*| = {error_norm:.5f}")

This simplified example showcases the power and flexibility of MLX, allowing researchers and developers to efficiently implement and train machine learning models on Apple silicon. Experimentation with MLX doesn’t stop here; it’s the beginning of a journey into a broader world of machine learning possibilities. Whether you’re optimizing existing models or innovating new ones, MLX offers the tools you need to make your mark in the field of AI research.

Join the MLX community

MLX’s journey is far from over; it thrives on community contribution. From improving documentation to extending the framework’s capabilities, every contribution helps shape MLX’s future.

MLX has already made a significant impact, thanks to its initial developers—Awni Hannun, Jagrit Digani, Angelos Katharopoulos, and Ronan Collobert. As more contributors join the fray, MLX’s potential only grows, promising to unlock new realms of possibility in machine learning research.

Conclusion

MLX represents a significant leap forward in making efficient and flexible machine learning research accessible on Apple silicon. Its innovative design, rooted in well-established frameworks, and bolstered by a unified memory model, paves the way for a new era of machine learning possibilities. Whether you’re a seasoned researcher or just setting foot in the world of machine learning, MLX extends an invitation to explore, innovate, and transform the landscape of AI research.

Explore MLX today and be part of the change driving the future of machine learning research on Apple silicon.