WWDC Quick Look 💓 By SwiftGGTeam
中文 EN 日本語
Explore numerical computing in Swift with MLX

Explore numerical computing in Swift with MLX

MLXSwiftNumerical ComputingAutoDiffGPU
观看原视频

Highlight

MLX Swift 把 NumPy 风格的 n 维数组运算和自动微分原生引入 Swift，让数值计算代码可以直接写成数学公式的样子，并自动在 GPU 上并行执行。

核心内容

数值计算在 Swift 里的痛点

做科学计算或机器学习的开发者，在 Apple 平台上长期面临一个尴尬选择：要么用 Accelerate 写底层向量运算，要么用 Metal Performance Shaders 手写 GPU kernel，要么干脆切换到 Python 生态。Swift 本身缺乏一套高层次的、表达力强的数值计算工具。写个矩阵乘法还能应付，一旦涉及物理模拟、迭代求解或模型训练，代码里就会充斥大量与数学无关的样板代码。

MLX Swift 的出现改变了这个局面。它的核心设计借鉴了 NumPy：以 n 维数组（MLXArray）为中心抽象，所有操作面向整个数组而非单个标量。如果你写过 NumPy，迁移到 MLX Swift 几乎不需要学习成本。

惰性求值：延迟到最后一刻

MLX Swift 的一个关键特性是惰性求值（lazy evaluation）。你对数组做的每一次加减乘除，并不会立即执行，而是构建一张计算图（compute graph）。直到你调用 eval() 或者读取某个具体值时，整个图才会被一次性提交给 GPU 执行。

这个设计的好处是操作融合（operator fusion）。GPU 启动 kernel 的开销很大，把多个小操作合并成一个大 kernel 能显著提升 Apple Silicon 上的性能。代价是开发者需要手动管理 eval() 的调用时机，尤其是在循环中，否则计算图会无限膨胀导致内存溢出。

自动微分：从 loss 到梯度，一行代码

训练模型的核心是梯度下降，而梯度下降的前提是算梯度。以前要在 Swift 里实现这个，要么手推导数公式，要么依赖 Swift for TensorFlow 那种编译器级别的深度改造。MLX Swift 通过函数变换 grad() 解决了这个问题：你把前向的 loss 函数丢进去，它返回一个自动计算梯度的函数。整个推导过程由 MLX 在底层完成，不需要你写一行导数代码。

详细内容

矩阵运算与幂迭代（03:04）

Session 以一个经典的幂迭代（power iteration）算法开场，演示 MLX Swift 的基本用法。目标是求一个对称矩阵的最大特征值和对应的特征向量。

import MLX
let n = 100
let steps = 10
let B = MLXRandom.normal([n, n])
var v = MLXRandom.normal([n])

// 构造对称矩阵 A = Bᵀ + B
let A = B.T + B

// 幂迭代：v ← A v / ‖A v‖
for _ in 0 ..< steps {
    let Av = matmul(A, v)
    v = Av / norm(Av)
    eval(v)  // 每步强制求值，防止计算图膨胀
}

// 恢复特征值：λ = vᵀ A v
let lambda = matmul(matmul(v.T, A), v)
print(lambda)

关键点：

  • B.T 取转置，+ 做矩阵加法，代码和数学公式一一对应
  • matmul 是矩阵乘法，norm 计算 L2 范数
  • 循环内必须调用 eval(v)，否则每次迭代都在往计算图里追加节点
  • 最后读取 lambda 时会触发惰性求值，自动完成所有未执行的运算

Mandelbrot 分形：从标量循环到数组运算（05:09）

Mandelbrot 集是展示数组计算威力的经典例子。对复平面上每个点 c，迭代公式 z = z² + c，如果模长始终不超过 2，该点属于集合。

纯 Swift 的标量实现需要双重嵌套循环，逐像素计算：

// 纯 Swift，逐标量计算
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
    }
}

MLX Swift 的版本把整个网格当作一个数组来处理：

import MLX

let x = linspace(-2.0, 0.5, count: w)
let y = linspace(-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                       // 对整个网格同时迭代
    counts = counts + (abs(z) .< 2)     // 统计未逃逸的迭代次数
}

关键点：

  • linspace 生成等间距序列，reshaped 调整维度以便广播
  • y.asImaginary() 把实数数组转成虚部，与 x 相加得到复数网格
  • z = z * z + c 一次性作用于整个数组，无需手动遍历每个像素
  • abs(z) .< 2 产生布尔掩码，统计每个点仍在边界内的迭代次数
  • 默认使用 GPU 执行，速度可比 CPU 标量版本快 10 倍以上

热传导模拟：卷积与边界条件（07:27）

热传导模拟展示了如何处理”每个单元格与邻居交互”的场景。Jacobi 迭代法的核心思想是：每个格子的新温度等于四个邻居的平均值。

// 卷积核：四个方向各 0.25，中心为 0
let kernel = MLXArray(converting: [
    0,    0.25, 0,
    0.25, 0,    0.25,
    0,    0.25, 0,
]).reshaped(1, 3, 3, 1)

var temperature = heatSources

// 循环直到收敛
let next = conv2d(temperature, kernel, padding: 1)
temperature = which(heatMask, heatSources, next)

关键点：

  • conv2d 一次性对整个网格应用卷积核，替代了手写邻居遍历
  • which 是元素级三元运算符：如果 heatMask 为真（热源或墙壁），保持固定值；否则取卷积结果
  • padding: 1 保证卷积后输出尺寸与输入一致

Jacobi 迭代虽然每步计算快，但收敛慢，热量每次只能传播一个格子。Session 接着介绍了逐次超松弛法（SOR），通过引入超松弛参数 omega 加速收敛：

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)

// 更新红色格子（用黑色邻居的旧值）
let sorRed = ω * conv2d(temperature, kernel, padding: 1) + (1 - ω) * temperature
temperature = which(redMask, sorRed, temperature)
temperature = which(heatMask, heatSources, temperature)

// 更新黑色格子（用红色邻居的新值）
let sorBlack = ω * conv2d(temperature, kernel, padding: 1) + (1 - ω) * temperature
temperature = which(blackMask, sorBlack, temperature)
temperature = which(heatMask, heatSources, temperature)

关键点：

  • omega 根据网格尺寸计算最优值，可将收敛迭代次数从 O(N²) 降到 O(N)
  • 红黑棋盘格掩码实现了”类原地更新”：红色格子更新时黑色邻居还是旧值，黑色格子更新时红色邻居已刷新
  • Session 演示中 SOR 比 Jacobi 快得多，甚至需要放慢 100 倍才能看清过程

曲线拟合与自动微分（11:13）

最后一个例子展示了 MLX Swift 的自动微分能力。给定一组数据点，用二次多项式拟合，核心代码如下：

// 定义预测函数
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     // 参数更新
    eval(θ)                      // 强制求值，截断计算图
}

关键点：

  • grad(loss) 把 loss 函数转换成返回梯度的函数，无需手写导数
  • gradLoss(θ) 在内部通过反向传播自动计算每个参数对 loss 的偏导
  • 循环末尾的 eval(θ) 同时完成计算和内存回收，漏写会导致 OOM
  • 这套机制是训练所有机器学习模型的基础，MLX Swift 把它简化到了几行代码

MLX 的完整工具箱（12:17）

Session 末尾列举了 MLX 的完整能力矩阵：线性代数、FFT、N 维卷积、归约操作、扫描前缀和、索引切片、随机数生成等。这些操作共享同一套惰性求值和自动微分基础设施。

核心启发

  • 端侧物理模拟 App：用 conv2d + which 组合实现流体、热传导或电场模拟。MLX Swift 的 GPU 加速让手机实时渲染复杂物理效果成为可能。入口 API：MLXArray、conv2d、which。

  • 自定义图像滤镜管线：利用数组级运算和 FFT 实现卷积神经网络风格的图像处理，比如风格迁移或超分辨率。MLX Swift 的惰性求值会自动融合多个滤镜操作为少量 GPU kernel。入口 API：MLXArray.fft、conv2d。

  • 端侧模型微调工具：基于 mlx-swift-lm 和 grad()，在设备上用小样本数据微调语言模型。MLX Swift 的自动微分让训练循环可以直接用 Swift 写，无需跨语言调用。入口 API：grad()、mlx-swift-lm 仓库。

  • 实时音频信号处理：用 FFT 和数组运算实现频谱分析、滤波器组或音频特效。MLX Swift 的 GPU 并行能力可以同时处理多路音频流。入口 API：MLXArray.fft、逐元素运算。

  • 交互式数学可视化：把 Mandelbrot 集扩展到 Julia 集、Newton 分形等，用 SwiftUI 做交互界面，MLX Swift 做后台计算。GPU 加速支持实时缩放和拖动。入口 API：linspace、reshaped、复数运算。

关联 Session

评论

GitHub Issues · utterances