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.
ml-exploreThe 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:
ml-explore- 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.
