使用 MLX 探索 Swift 中的数值计算

使用 MLX 探索 Swift 中的数值计算

使用 MLX Swift 将 NumPy 风格的计算原生引入 Swift。了解如何在一个类型安全的环境中统筹管理图像处理、张量运算和神经网络训练等任务，从而消除机器学习工作流程中的跨语言障碍。探索相关 API，尽享熟悉的编译器、工具和调试体验，同时充分发挥 GPU 加速的优势。

章节
- 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
资源
相关视频

WWDC26
- 使用 MLX 在 Mac 上本地运行 AI 智能体
- 使用 MLX 探索分布式推理和训练
WWDC25
- 借助 MLX 在 Apple 芯片上探索大语言模型
- 开始使用适用于 Apple 芯片的 MLX
你好我是David Koski 我在MLX Swift团队工作数值计算也称为数值分析或科学计算是一套用于解决数学问题的技术和算法这些问题往往难以通过符号或手工求解它们需要大量的计算应用领域包括化学仿真生物学物理学和金融系统其他领域包括音频和信号处理视觉应用包括渲染光线追踪和分形大规模梯度下降可以进行任意曲线拟合这是训练机器学习模型的基础今天我们将讨论使用MLX Swift进行数值计算 Apple平台拥有丰富的数值计算生态系统每个现有框架在其设计领域都表现出色 Accelerate为你提供 CPU上经过手动调优的向量原语 BNNS是神经网络的基础构建块层 Metal Performance Shaders 让你直接访问GPU内核 Swift Numerics添加了复数类型和泛型数值协议那么何时使用MLX Swift呢如果你的主要目标是编写数学代码同时注重性能 MLX Swift是一个很好的解决方案你所编写的代码看起来就像你实现的数学公式无需底层库的编程开销也不需要在纯Swift中操作数组时繁琐的记录工作它为何看起来像数学核心思想很简单数学家和数值分析师使用向量和矩阵进行运算不是对单个值执行运算而是一次性对整个矩阵进行操作 MLX Swift以n维数组作为核心抽象就像NumPy和许多其他框架一样事实上如果你用过NumPy 这个API会非常熟悉大多数NumPy代码只需少量修改即可转换为MLX Swift 这种表达方式极具表现力且简洁明了你编写时能理解之后阅读时也能看懂数组计算和惰性求值使两件事成为可能自动GPU执行和自动微分最棒的是 mlx-swift 和整个MLX生态系统都是基于MIT许可证的开源项目我们欢迎提交Issue和PR 社区非常活跃欢迎提问修复Bug 让它变得更好以下是今天的计划我将先介绍MLX Swift 以及一些基本的矩阵向量运算然后深入三个示例展示 MLX Swift如何轻松将数学转化为代码 Mandelbrot集的计算求热分布的稳态以及最后的曲线拟合首先我们来谈谈MLX Swift 这里展示了一些MLX Swift操作以幂迭代法为例让我们逐步讲解我们导入MLX 设置矩阵大小和迭代次数然后从正态分布中采样随机矩阵和向量接下来通过将B与其转置相加构建对称矩阵注意这与数学公式有多相近 .T得到转置加号执行矩阵加法在循环内部 matmul执行矩阵向量乘法 norm给出L2范数同样代码读起来就像数学这里我们还看到MLX的一个关键特性：惰性求值对MLX数组对象的操作会构建一个计算图只有当你调用eval 或读取值时才会执行在这样的循环中每步调用eval以保持计算图小巧惰性求值也是支持 MLX函数变换的基础例如用于自动微分的grad
最后我们恢复特征值读取值会触发计算的执行与其他数组框架一样 MLX Swift代码读起来几乎像数学公式如果你实际上需要矩阵的所有特征值和特征向量 MLX Swift线性代数包也提供了相应函数让我们继续下一个示例接下来是Mandelbrot集它是一个经典分形也是数组计算的完美展示案例我们对大量网格点应用一个函数其定义出人意料地简单对于复平面上的每个点c 迭代 z = z² + c 如果其模从未超过2 该点在集合内并被染成黑色如果发散则根据逃逸速度着色分形之美在于它是自相似的且无限精细你可以永远平移和缩放图案永不重复我们先从使用标量的纯Swift实现开始你遍历每个像素运行Mandelbrot迭代并检查是否发散它能正常工作是符合Swift惯例的代码但你需要管理大量与问题本身无关的记录工作而且它在CPU上运行每次处理一个点让我们来看MLX Swift
这是MLX Swift的实现设置一个复数网格即c 然后循环只需两行代码 z = z * z + c 同时应用于所有点统计每个点保持有界的迭代次数就这样代码是数学公式的直接翻译计算在整个网格点上执行就像对单个点操作一样简单默认使用GPU 带来快速的性能
纯Swift非常具有表现力你可以自然地编写数值计算代码但你每次只处理一个标量所以你需要自己遍历每个点繁琐的记录工作会让数学逻辑变得模糊
MLX Swift专为数值计算而构建你对数组而非标量进行操作代码看起来就像你想表达的数学公式它在GPU上运行更快并行处理所有点快多少取决于具体算法但快10倍绝对是可能的这一切都用更少更简洁的代码实现 Mandelbrot是典型的并行问题每个点都是独立的下一个示例不同每个单元格与邻居交互这种模式在物理学图像处理和神经网络中随处可见 MLX用一个操作处理它即卷积想象一个有墙壁和热源的房间我们想知道内部各处的稳态温度求解这个问题的最简单方法称为Jacobi迭代法将温度建模为二维网格
每次新迭代对相邻值取平均使用这样的模板你不断重复此过程热量向外扩散直到达到稳态注意更新只关注一个小邻域而且在每个点使用相同的规则这正是卷积的本质让我们看看代码这是MLX Swift中求解器的核心代码让我来逐步讲解卷积核就是上一张幻灯片的模板直接写出来四个邻居各占四分之一中心和角落为零温度网格从热源开始初始化这是一个合理的初始值循环内部只有两行第一行是物理部分 conv2d将模板一次性应用于整个网格第二行处理边界条件这是一个逐元素的三元运算凡是掩码标记为热源或墙壁的地方保留固定值其他地方取卷积得到的新值就这样数学说对四个邻居取平均我们用一次conv2d调用实现了这一点 Jacobi迭代计算速度快但收敛缓慢热量每次只移动一个单元格通常达到稳态需要N²次迭代其中N是网格的边长就像快速排序比冒泡排序做更少的工作有些算法能用更少的工作达到稳态其中一种称为逐次超松弛法（SOR）该方程看起来类似于Jacobi迭代实际上它使用相同的卷积核它使用参数omega 将每次更新在变化方向上推得更远稍微过冲以更快地到达目标过冲会在迭代中逐渐恢复 omega参数可以根据数组的大小计算得出使用最优值时将在N次迭代内收敛该技术的另一个关键是原地更新 MLX通常生成新数组而不是原地更新但红黑棋盘格模式通过交替处理单元格可用于计算新值达到相同的效果来看代码首先根据网格大小计算最优omega 你可以看到我们在这里如何使用omega 它与方程完全一致我们使用棋盘格掩码交替更新数组中的单元格现在黑色单元格运行相同的更新但这次它们的红色邻居已经有了最新的值这正是我们需要的原地更新效果像之前一样在循环中重复让我们看看差异 Jacobi和SOR的代码几乎完全相同而且都与数学公式高度吻合让我们看看它们的运行效果你可以看到上方的Jacobi 缓慢扩散而下方的SOR 快速填满区域 SOR运行时呈现出明显的波纹图案那是实时的过冲和修正过程最终两者都收敛到相同的状态还有一件事我不得不将SOR放慢100倍才能让它清晰可见选择正确算法的威力前两个示例是正向计算从输入开始计算输出最后一个则相反你有输出即数据点你想找到能产生这些数据的参数这是我们将应用一个关键函数变换的地方由MLX Swift提供用于自动微分的grad 假设你有一些数据点你想找到一个函数来近似这些点你决定所需函数的结构可以是多项式正弦之和或任何你喜欢的形式在这个例子中我们将使用多项式一个给我们抛物线的二次函数你想最小化损失即函数输出与实际数据之间差值平方的均值这与训练每个机器学习模型的核心思想相同只是规模更小
要最小化损失我们需要关于参数的梯度以及用于更新参数的优化循环我们定义函数f 和损失函数即均方误差接下来我们创建theta 即要拟合的系数并将损失函数变换为一个能产生精确梯度的函数关于参数的梯度我们无需手动编写任何导数这些都是由MLX推导出来的在优化循环中我们在当前参数处计算梯度迈出一小步并调用eval 清空计算图每次迭代都如此防止计算图无限增长这就是梯度下降
这是效果展示抛物线很快接近并超过数据然后逐渐收敛到越来越接近的近似值这个示例用的是简单多项式我们本可以用线性代数包中的QR 直接拟合曲线梯度适用于任意复杂的函数如果你需要的不仅仅是梯度 MLX提供了一套优化算法如SGD Adam RMSprop等
今天我展示了数组计算卷积和grad 但MLX包含完整的数值计算工具包以下是部分示例线性代数 FFT N维卷积归约运算扫描运算索引随机数生成以及更多功能基于MLX Swift构建的软件包生态系统已相当健全核心mlx-swift仓库就是本次Session中展示的框架 mlx-swift-lm是Swift语言模型实现所在的地方 mlx-swift-examples包含可供入门参考的示例程序本次Session的示例也将发布在那里所有这些都是开源的只需几行代码就可通过Swift Package Manager安装 MLX不只支持Swift 它是一个具有四个前端的框架 Swift Python C++和C 第三方还构建了更多前端以满足更多需求它们共享相同的概念相同的操作以及相同的惰性求值模型概念和模式在它们之间只需少量修改即可迁移因此你可以用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).

探索“入门汇总”

及时了解最新动态

探索“平台”

精选

探索“技术”

精选

探索“社区”

精选

探索“文档”

发布说明

探索“下载”

精选

探索“支持”

精选

快速链接

章节

资源

相关视频

WWDC26

WWDC25