Swift Digest
この記事の要点
- Apple silicon 上で機械学習のリサーチを行うための配列フレームワーク MLX を Swift から使えるようにする MLX Swift が発表されました。MIT ライセンスで公開されています。
- MLX Swift には、MLX コアの包括的な Swift API、ニューラルネットワークと最適化アルゴリズムの高レベルなパッケージ、Mistral 7B によるテキスト生成や MNIST 学習のサンプル、Swift と C++ コアをつなぐ C API が含まれます。
- ハードウェアアクセラレーション（CPU・GPU での演算）と、ニューラルネットワークの学習に使う自動微分（automatic differentiation）に対応しており、既存の Swift ライブラリではほとんど扱えなかった機械学習リサーチ向けの機能を提供します。
- MLX はリサーチ向けであり、アプリでのモデルの本番デプロイを想定したものではありません。
何が発表されたのか
Swift は、Python のような言語の使いやすさと高水準な構文を、C++ のようなコンパイル言語の速度と兼ね備えており、機械学習のリサーチに使う言語として大きな可能性があります。
MLX は、Apple silicon 上での機械学習リサーチのための配列フレームワークです。MLX Swift は、その MLX を Swift から扱えるように拡張するもので、機械学習の研究者が Apple silicon 上で手軽に実験できるようにします。
今回のリリースには次が含まれます。
- MLX コアの包括的な Swift API
- 高レベルなニューラルネットワークと最適化アルゴリズム（optimizers）のパッケージ
- Mistral 7B によるテキスト生成のサンプル
- MNIST 学習のサンプル
- Swift と C++ コアの橋渡しをする MLX への C API
これらはすべて MIT ライセンスのもとで公開されています。
MLX には、既存の Swift ライブラリではほとんど対応されていなかった、機械学習リサーチにとって重要な機能が備わっています。
- ハードウェアアクセラレーションへのネイティブ対応。計算負荷の高い演算を CPU でも GPU でも実行できます。
- ニューラルネットワークや勾配ベースの機械学習モデルを学習するための自動微分（automatic differentiation）。
何に使えるのか
MLX Swift は Xcode や SwiftPM から手軽にセットアップでき、N 次元配列の構築と演算をシンプルに書けます。次の例では、特に指定しない限りデフォルトのデバイス（GPU）上ですべての演算が実行されます。
import MLX
import MLXRandom
let r = MLXRandom.normal([2])
print(r)
// array([-0.125875, 0.264235], dtype=float32)
let a = MLXArray(0 ..< 6, [3, 2])
print(a)
// array([[0, 1],
// [2, 3],
// [4, 5]], dtype=int32)
// last element of 0th row
print(a[0, -1])
// array(1, dtype=int32)
// slice of the first two rows
print(a[0 ..< 2])
// array([[0, 1],
// [2, 3]], dtype=int32)
// add with broadcast
let b = a + r
print(b)
// array([[-0.510713, 1.04633],
// [1.48929, 3.04633],
// [3.48929, 5.04633]], dtype=float32)
MLX Swift では関数変換（function transformation）も利用できます。これはモデルを自動微分で学習させたり、計算グラフを速度やメモリ使用量の面で最適化したりするのに役立ちます。次の例は、ある関数の勾配を計算します。
func fn(_ x: MLXArray) -> MLXArray {
x.square()
}
let gradFn = grad(fn)
let x = MLXArray(1.5)
let dfdx = gradFn(x)
// prints 2 * 1.5 = 3
print(dfdx)
ドキュメントには、より完全なサンプルも用意されています。
- LLM によるテキスト生成: Mistral 7B を使った完全なテキスト生成の例です。あらかじめ量子化された MLX モデルを含む、Mistral 系・Llama 系のモデルでテキストを生成できます。これらのモデルの多くは Hugging Face から入手できます。
- MNIST での MLP の学習: MLX Swift のニューラルネットワークと最適化アルゴリズムのパッケージを使って、MNIST の数字を分類する単純な多層パーセプトロンを学習する例です。
導入・今後の位置づけ
MLX Swift は、機械学習の研究者がさまざまなプラットフォームやデバイスで手軽に実験できる、研究者にとって扱いやすい機械学習フレームワークです。導入を始めるには、Swift 向けのドキュメントとサンプルや GitHub リポジトリが参考になります。問題を見つけたり改善の提案があれば issue を、貢献したい場合はコントリビューションガイドラインを通じて参加できます。
なお、MLX はあくまでリサーチ向けであり、アプリでのモデルの本番デプロイを想定したものではない点に注意が必要です。