-
MLXを活用したSwiftでの数値計算
MLX Swiftを使うと、NumPyスタイルのコンピューティングをSwiftでネイティブに実行できます。画像処理、テンソル演算、ニューラルネットワークのトレーニングを単一の型安全な環境で実行することで、機械学習ワークフローにおける言語間の摩擦を解消する方法を確認しましょう。使い慣れたコンパイラ、ツール、デバッグの利用体験を変えることなく、GPUアクセラレーションを活用できるようにするAPIも紹介します。
関連する章
- 0:00 - Introduction
- 0:57 - MLX Swift and the Apple ecosystem
- 3:04 - MLX Swift
- 4:28 - Mandelbrot
- 6:34 - Heat distribution
- 8:12 - Faster convergence with SOR
- 10:17 - Curve fitting
- 12:17 - The full MLX toolkit and ecosystem
- 13:47 - Next steps
リソース
- MLX Swift LM on GitHub
- MLX Swift Examples
- MLX Examples
- MLX Swift
- MLX LM - Python API
- MLX Explore - Python API
- MLX Framework
- MLX
関連ビデオ
WWDC26
WWDC25
-
このビデオを検索
こんにちは David Koskiです MLX Swiftを担当しています 数値計算は数値解析とも 科学計算とも呼ばれます 数学的問題を解くための 手法やアルゴリズムの集合です これらは多くの場合 記号的や手動では解くのが難しいです 膨大な計算量を必要とします 用途には化学の シミュレーションがあります 生物学 物理 金融システムも 含まれます 他の分野には音声や 信号処理もあります 視覚的な用途にはレンダリング レイトレーシング フラクタルがあります 大規模な勾配降下法は 任意のカーブフィッティングが可能です これが機械学習モデルの 訓練の基礎となります 今日はMLX Swiftを使った 数値計算について説明します Appleプラットフォームには 豊富な数値計算エコシステムがあります 各フレームワークはその 設計目的に優れています AccelerateはCPUでの 手調整ベクタープリミティブを提供します BNNSはニューラルネットワークの 基本ブロック層です Metal Performance Shadersは GPUカーネルに直接アクセスできます Swift NumericsはComplex型と 汎用数値プロトコルを追加します ではMLX Swiftはいつ使うのでしょうか 主な目的が数学的なコードを書くことで パフォーマンスを重視する場合 MLX Swiftは優れた選択肢です 書くコードは実装している数学に似ており 低水準ライブラリのような プログラミングオーバーヘッドがなく 通常のSwiftで配列を操作する際の 詳細な管理も不要です どのように数学のように見えるのでしょうか 核となるアイデアはシンプルです 数学者や数値解析者は ベクターや行列を扱います 単一の値に演算を行うのではなく 行列全体に一度に演算を行います MLX Swiftはn次元配列を 中心的な抽象化として使います NumPyや多くの フレームワークと同様に NumPyを使ったことがあれば APIは非常に見覚えがあります ほとんどのNumPyコードは 最小限の変更でMLX Swiftに変換できます 非常に表現力が高く 複雑ではありません 書いた時も後で読んだ時も コードを理解できます 配列計算と遅延評価により 2つのことが可能になります GPUの自動実行と 自動微分です 何よりmlx-swiftと MLXエコシステム全体は すべてMITライセンスの オープンソースです issueやPRを歓迎しており 活発なコミュニティがあります 質問し バグを修正し より良くしましょう！ 今日のプランを説明します まずMLX Swiftの基本的な 行列ベクター演算を紹介します 次に3つの例を掘り下げます MLX Swiftが数学を コードに変換しやすいことを示します マンデルブロ集合の計算 熱分布の定常状態の求め方 そして最後にカーブフィッティングです まずMLX Swiftについて話しましょう いくつかのMLX Swift演算を示します べき乗法を例として使います 順を追って見ていきましょう MLXをインポートし 行列サイズと反復回数を設定します 次に正規分布からランダムな 行列とベクターをサンプリングします BとBの転置を加算して 対称行列を作ります これが数学にいかに近いかに 注目してください .Tが転置を与え プラスが行列加算を行います ループ内でmatmulが 行列ベクター積を行います normがL2ノルムを返します コードが数学のように読めます MLXの重要な機能である 遅延評価も確認できます MLX配列オブジェクトへの演算は 計算グラフを構築します evalを呼ぶか値を読み取るまで 何も実行されません このようなループでは各ステップで evalを呼びグラフを小さく保ちます 遅延評価はMLXの 関数変換も支えています 自動微分のためのgradなどです
最後に固有値を取り出し 値を読み取ると計算が強制されます 他の配列フレームワークと同様に MLX Swiftのコードは 数学のように読めます すべての固有値が実際に必要なら 行列の固有ベクターも含め MLX Swift線形代数パッケージに それらの関数もあります 次の例に進みましょう 次はマンデルブロ集合です 古典的なフラクタルで 配列計算の完璧な例でもあり 大きなポイントグリッドに 関数を適用します 定義は驚くほどシンプルです 複素平面上のすべての点cに対して z = z² + cを反復します 大きさが2を超えなければ その点は集合内にあり黒く塗られます 発散する場合は 逃げる速さで色付けされます フラクタルの美しさは 自己相似性と無限の細部にあります ずっとパンやズームができ パターンは繰り返されません スカラーを使った通常のSwiftの 実装から始めましょう すべてのピクセルをループして マンデルブロの反復を実行し 発散をチェックします 動作します Swiftらしいコードです しかし問題とは無関係な 多くの管理が必要です CPUで1点ずつ実行されます MLX Swiftを見てみましょう
これがMLX Swiftでの実装です 複素数のグリッドを設定します これがcです ループはたった2行です z = z * z + c すべてのポイントに一度に適用されます 各ポイントが有界に留まる 反復回数を数えます 以上です コードは数学の直接的な変換です 計算はポイントの グリッド全体に対して行われます 単一ポイントと同じくらい簡単に デフォルトではGPUが使われ 高速なパフォーマンスが得られます
通常のSwiftは表現力があります 数値計算コードを自然に書けます しかし1スカラーずつ処理となり すべてのポイントを自分で ループする必要があります 管理コードが数学を覆い隠してしまいます
MLX Swiftは数値計算向けに 構築されており スカラーではなく配列で演算します 表現しようとしている数学のように見えます GPUでより速く動作し すべてのポイントを並列処理します 速度向上は使用する アルゴリズムによって異なりますが 10倍の速度向上も十分に可能です これらすべてがより小さく シンプルなコードで実現できます マンデルブロは自明なほど並列でした すべてのポイントが独立しています 次の例は異なります 各セルが隣接セルと連携します そのパターンは物理 画像処理 ニューラルネットワークなどに現れます MLXは1つの演算である 畳み込みで処理します 壁と熱源がある部屋を想像してください 内部のすべての場所での 定常温度を求めたいです これを解く最も簡単な方法は Jacobi法です 温度を2Dグリッドでモデル化します
各新しい反復は隣接値の平均を取ります このようなステンシルを使って これを繰り返すと熱が広がっていきます 定常状態に達するまで 更新は小さな近傍だけを参照します すべてのポイントで同じレシピを使います これはまさに畳み込みそのものです コードで見てみましょう MLX Swiftのソルバのコアです 順を追って見ていきましょう カーネルは前のスライドのステンシルで 文字通りに書きます 隣接する4セルに0.25 中心と角には0を指定します 温度グリッドは熱源の値から始まります 合理的な初期値です ループの中は2行です 1つ目は物理演算です conv2dは1回の呼び出しで グリッド全体にステンシルを適用します 2行目は境界条件を処理します 要素ごとの三項演算子です マスクが熱源や壁と示す箇所は 固定値を保持し それ以外は畳み込みの新しい値を使います 以上です 数学では4隣接セルの平均を取ります それをconv2dへの1回の 呼び出しで実装しました Jacobi法は計算は速いが 収束が遅いです 熱は1セルずつしか移動できません 通常 定常状態には N^2回の反復が必要で Nはグリッドのサイズのことです クイックソートが バブルソートより効率的なように 定常状態により少ない計算量で 到達できるアルゴリズムがあります その1つはSOR 逐次過緩和法と呼ばれます 数式はJacobi法に似ています 実は同じ畳み込みカーネルを使います ωというパラメータを使い 各更新を変化の方向に さらに押し進め わずかにオーバーシュートして速く到達します 反復するにつれて オーバーシュートは収束します ωパラメータは配列サイズに基づいて 計算できます 最適値を使うと N回の反復で収束します この手法のもう1つの鍵は インプレース更新です MLXは通常インプレース更新ではなく 新しい配列を生成しますが 赤と黒のチェッカーボードパターンを 交互のセルを処理することで 新しい値を計算できます 同じ効果が得られます コードを見ていきましょう！ まずグリッドサイズに基づいて 最適なωを計算します ここでωの使い方を確認できます 数式と完全に一致しています チェッカーボードマスクを使って 配列の交互のセルを更新します 次に黒いセルが同じ更新を実行します 今回は赤い隣接セルが すでに新しい値を持っています これがまさに必要としていた インプレース効果です これを前と同様ループで繰り返します 違いを確認しましょう JacobiとSORのコードはほぼ同一で 数学と密接に一致しています 実行結果を確認しましょう 上にJacobiが見えます ゆっくり広がり 下のSORは素早くエリアを埋めていきます SORは実行中に 顕著な波紋パターンを示します リアルタイムでの オーバーシュートと修正です 最終的に両方が 同じ状態に収束します もう1点あります SORを表示できるよう 100倍遅くしました 適切なアルゴリズムを選ぶ力です！
最初の2つの例は 順方向の計算でした 入力から出力を計算します 最後の例はそれを逆にします 出力とデータポイントがあり それらを生成するパラメータを求めたいです ここで重要な関数変換を適用します MLX Swiftが提供する 自動微分のためのgradです いくつかの点があるとします それらに近似する関数を見つけたいです 使いたい関数の構造を決めます 多項式 正弦波の和など 好きなものを選べます この例では多項式を使います 放物線を与える二次式です 損失を最小化したいです 関数の出力との差の二乗平均 と実際のデータです これはすべてのMLモデルの 訓練の背後にある核心と同じで ただし小さなスケールです
損失を最小化するには パラメータに関する勾配が必要で パラメータを更新する 最適化ループも必要です 関数fを定義します と損失を定義し 平均二乗誤差を使います 次にフィットする係数θを作り 損失関数を変換して 正確な勾配を生成する関数に変えます パラメータに関して 微分を手で書く必要はありません MLXが導出しました 最適化ループでは 現在のパラメータで勾配を評価し 小さなステップを取りevalを呼んで 計算グラフをフラッシュします 各反復で無限に大きくならないようにします これが勾配降下法です
実際にどのように見えるか確認しましょう 放物線はすぐにデータに近づき オーバーシュートします しかしより近い近似へと落ち着いていきます この例は単純な多項式でした 線形代数パッケージのQRを使って 直接カーブをフィットすることもできました 勾配は任意の複雑な関数に対応します 勾配以上のものが必要な場合 MLXはSGDなどの 最適化アルゴリズムを豊富に揃えています Adam RMSpropなどがあります
今日は配列計算 畳み込み gradを紹介しましたが MLXは完全な数値計算 ツールキットを備えています サンプルをご覧ください 線形代数 FFT N次元畳み込み リダクション スキャン インデキシング 乱数生成などがあります MLX Swiftの上に構築された パッケージの豊富なエコシステムがあります 中心となるmlx-swiftリポジトリは このセッションのフレームワークです mlx-swift-lmにはSwift 言語モデルの実装があります mlx-swift-examplesには 入門用のサンプルプログラムがあります このセッションを元にした例も そこに投稿されます これらはすべてオープンソースで数行で Swift Package Managerで導入できます MLXはSwiftだけではありません Swift Python C++ C フロントエンドは4つあります 必要に応じてサードパーティが さらにフロントエンドを作っています 同じコンセプトを共有し 同じ演算を共有し 同じ遅延評価モデルを共有します 概念とパターンは 最小限の変更で引き継がれます PythonでプロトタイプしSwiftで出荷できます Pythonの周りにはより広い 研究エコシステムがあります mlx-lmやmlx-vlmは注目に値します Python側で何が作られているか 確認したいなら mlx-swiftとmlx-swift-examples をぜひご覧ください mlx-swiftにはドキュメントがあります 動作を確認できるテストもあります mlx-swift-examplesには 様々なサンプルアプリがあります LLM統合を実演するものもあります Stable Diffusion モデルのトレーニングと ファインチューニング そしてもちろんこのトークで 紹介した例もあります 数値計算のニーズがありますか あるいはただ試してみたい 面白いシミュレーションや視覚化で ぜひ試してみてください！ 何かひらめいたら 皆さんに参加と貢献を呼びかけます 修正できるオープンなissueもあります または新しいサンプルを プログラムとして作れます ご視聴ありがとうございました
-
-
3:04 - Power iteration with MLX Swift arrays
import MLX let n = 100 let steps = 10 let B = MLXRandom.normal([n, n]) var v = MLXRandom.normal([n]) // get symmetric matrix A = Bᵀ + B let A = B.T + B // Power iteration → top eigenvector of A. // v ← A v / ‖A v‖ for _ in 0 ..< steps { let Av = matmul(A, v) v = Av / norm(Av) eval(v) } // recover the eigenvalue. // λ = vᵀ A v let lambda = matmul(matmul(v.T, A), v) print(lambda) -
5:09 - Mandelbrot set in plain Swift (scalar)
// Plain Swift, scalar-at-a-time var counts = Array2D<Int>(width: w, height: h) for y in 0 ..< h { for x in 0 ..< w { let c = Complex(xMin + Float(x) * xStep, yMin + Float(y) * yStep) var z = Complex<Float>.zero var limit = maxIterations for i in 0 ..< maxIterations { z = z * z + c if z.lengthSquared > radiusSquared { limit = i break } } counts[x, y] = limit } } -
5:27 - Mandelbrot set in MLX Swift (array)
// Compute the Mandelbrot set on a grid of complex numbers import MLX let x = linspace(Float(-2.0), 0.5, count: w) let y = linspace(Float(-1.25), 1.25, count: h).reshaped(h, 1) let c = x + y.asImaginary() var z = MLXArray.zeros(like: c) var counts = MLXArray.zeros(c.shape, dtype: .int16) for _ in 0 ..< maxIterations { z = z * z + c // iterate z ← z² + c counts = counts + (abs(z) .< 2) // count bounded iterations } -
7:27 - Jacobi iteration with conv2d
// Jacobi iteration: average the four neighbors // Convolution weights let kernel = MLXArray(converting: [ 0, 0.25, 0, 0.25, 0, 0.25, 0, 0.25, 0, ]).reshaped(1, 3, 3, 1) // Initial value var temperature = heatSources // Run this in a loop until convergence let next = conv2d(temperature, kernel, padding: 1) temperature = which(heatMask, heatSources, next) -
9:17 - Successive Over-Relaxation (SOR)
// Successive Over-Relaxation: blend the previous and next state let ω: Float = 2.0 / (1.0 + sin(Float.pi / Float(max(M, N)))) let redMask = checkerboard(rows: M, cols: N, phase: 0) let blackMask = checkerboard(rows: M, cols: N, phase: 1) // Update red cells using black neighbors let sorRed = ω * conv2d(temperature, kernel, padding: 1) + (1 - ω) * temperature temperature = which(redMask, sorRed, temperature) temperature = which(heatMask, heatSources, temperature) // Update black cells using (now-updated) red neighbors let sorBlack = ω * conv2d(temperature, kernel, padding: 1) + (1 - ω) * temperature temperature = which(blackMask, sorBlack, temperature) temperature = which(heatMask, heatSources, temperature) -
11:13 - Curve fitting with automatic differentiation
// Define a loss, then optimize it with autodiff // x, y: data points as MLXArrays func f(_ θ: MLXArray) -> MLXArray { θ[0] + θ[1] * x + θ[2] * x ** 2 } func loss(_ θ: MLXArray) -> MLXArray { mean((f(θ) - y) ** 2) } var θ = zeros([numParams]) let gradLoss = grad(loss) for _ in 0 ..< steps { let g = gradLoss(θ) // ∇L(θ) θ = θ - learningRate * g // parameter update eval(θ) // force evaluation }
-
-
- 0:00 - Introduction
What numerical computing is and its applications — from simulations in chemistry, biology, and physics to signal processing, rendering, fractals, and machine learning training via gradient descent.
- 0:57 - MLX Swift and the Apple ecosystem
Where MLX Swift fits among Apple's existing numerical computing frameworks (Accelerate, BNNS, Metal Performance Shaders, Swift Numerics) — and why to choose MLX Swift when your primary goal is writing mathematical code that looks like the math, with automatic GPU execution and automatic differentiation.
- 3:04 - MLX Swift
Core MLX Swift concepts — n-dimensional arrays as the central abstraction (similar to NumPy), lazy evaluation that builds a compute graph before executing, and a walkthrough of the power iteration algorithm showing how operations like matmul, norm, and transpose map directly to mathematical notation.
- 4:28 - Mandelbrot
Computing the Mandelbrot fractal as a showcase for array computing — comparing a scalar-at-a-time plain Swift implementation against an MLX Swift version that applies z = z² + c across the entire grid of complex numbers at once, running on the GPU with up to 10x speedup in fewer lines of code.
- 6:34 - Heat distribution
Finding steady-state temperature in a room using Jacobi iteration — modeling heat as a 2D grid, implementing the four-neighbor stencil as a single conv2d call, and applying boundary conditions with an elementwise ternary. A convolution as physics.
- 8:12 - Faster convergence with SOR
How Successive Over-Relaxation (SOR) reaches steady state in N iterations versus N² for Jacobi — using an omega parameter to overshoot updates, a red/black checkerboard pattern to simulate in-place updates, and a side-by-side comparison showing SOR completing 100x faster with the same minimal code.
- 10:17 - Curve fitting
How automatic differentiation flips forward computation — given data points and a parametric function (a quadratic), using MLX's grad to derive gradients automatically and running a gradient descent loop to fit the curve without writing a single derivative by hand.
- 12:17 - The full MLX toolkit and ecosystem
Overview of MLX's complete numerical computing toolkit — linear algebra, FFTs, n-dimensional convolutions, reductions, scans, indexing, random number generation, optimizers (SGD, Adam, RMSprop) — and the Swift ecosystem of packages including mlx-swift, mlx-swift-lm, and mlx-swift-examples.
- 13:47 - Next steps
How to get started with mlx-swift and mlx-swift-examples (LLM integration, stable diffusion, model training, and session examples), MLX's multi-language support (Swift, Python, C++, C).