Apple Joins the AI Race with Release of MLX, A New Framework for Machine Learning on Apple Silicon

Apple unveils MLX, an ML framework crafted for Apple silicon. Early benchmarks show promise of 30% speedups over PyTorch, but can it handle training state-of-the-art models on consumer Macs?

Apple Joins the AI Race with Release of MLX, A New Framework for Machine Learning on Apple Silicon

Earlier this month, Apple made the surprise announcement that its machine learning research team has released MLX - a new open-source machine learning framework designed specifically for Apple silicon. This unexpected move signals Apple's intention to compete in the rapidly advancing AI space.

GitHub - ml-explore/mlx: MLX: An array framework for Apple silicon
MLX: An array framework for Apple silicon. Contribute to ml-explore/mlx development by creating an account on GitHub.

The Goal Behind MLX

Unlike Apple's typical secrecy around proprietary software, the company chose to make MLX open source. Their aim is to provide developers a specialized framework to build models that leverage the AI capabilities of Apple silicon to run efficiently.

As engineer Ani Hanon tweeted, MLX will allow users to "train a Transformer LM or fine-tune Léon on it. You can do text generation with LLMs, image generation with Stable Diffusion and speech recognition with Whisper." This demonstrates the wide range of cutting-edge AI applications enabled by the new framework.

Key features of MLX

Some of the key features that set MLX apart include:

  • Familiar APIs
    The Python API closely mirrors NumPy, allowing easy integration with existing workflows. On the C++ side, naming conventions match the Python side for straightforward cross-language development.
  • Higher-level Packages
    Convenience packages like mlx.nn and mlx.optimizers provide PyTorch-style abstractions to simplify model building on top of the base arrays.
  • Function Transformations
    Built-in transformations handle differentiation, vectorization, and graph optimizations under the hood to accelerate model training.
  • Lazy computation
    Computations are delayed until values are explicitly needed, avoiding unnecessary work and allowing dynamic graph construction.
  • Multi-Device Support
    MLX can target CPU, GPU, or any other supported devices within a single process using a unified memory model. No copies needed!

The goal of MLX is to produce an easy-to-use yet performant framework tailored for machine learning workloads on Apple hardware. The simple conceptual model takes inspiration from NumPy, PyTorch, Jax, and ArrayFire while targeting Apple's silicon with novel techniques.

Early Benchmarks Look Promising

Some early adopters have already started benchmarking MLX on available Apple silicon hardware. While full models haven't yet been ported, initial small-scale tests indicate MLX achieves up to 30% better performance over PyTorch for some workloads. As the software matures and is better optimized for the hardware, gains may become even more significant.

Initial reception:

  • Promising idea with potential to grow.
  • Currently runs Llama and Mistral at peak performance.
  • Needs further development and optimization.

Potential use cases:

GitHub - ml-explore/mlx-examples: Examples in the MLX framework
Examples in the MLX framework. Contribute to ml-explore/mlx-examples development by creating an account on GitHub.
  • Training large language models on Apple hardware.
  • Running computationally intensive ML tasks on Apple devices.

The Great Debate: Training vs Inference

There is some debate around whether MLX will work best for training models from scratch or fine-tuning and inference tasks. Smaller models can certainly train effectively on Apple silicon hardware. However, larger state-of-the-art models may still require additional compute power. Time will tell where the limits truly lie once more models are ported over!

Early Signs of an Emerging Apple AI Foundation Model?

The launch of MLX lays the groundwork for Apple to develop its own foundation model to rival models like GPT from Anthropic and Claude from Anthropic. While the company remains tight-lipped about future plans, providing developers these tools could lead to an ecosystem of MLX models.

So far, Apple has open sourced MLX examples for tasks like text generation, image generation, and speech recognition. My own experiments confirm MLX's potential - training a small 50 million parameter LM on an M2 Max utilized the full GPU with decent performance. More complex experiments are still needed, but it's an exciting first step.

With increasing competition in AI across the industry, Apple's release of this specialized framework signals its ambitions to compete using the differentiation of its silicon. By open sourcing MLX, it could also spur innovation from its developer community to push Apple's AI capabilities forward. This launch kicks off the next phase in Apple's integration of AI into its products and services.

Closing Thoughts

While MLX is still young, it shows exciting potential as an Apple-centric machine learning framework. For Mac diehards, this may finally open the door for effective ML development on Apple hardware. If performance claims hold up, MLX and future Apple silicon chips may give NVIDIA a run for its money! We'll be keeping a close eye on new developments as the library evolves.

Read next