We rewrote an ML Framework* in Lean, (and yes it is faster*)
This article describes the experience of rewriting tinygrad, a lightweight machine learning framework, in the Lean programming language. The authors discuss the motivations behind the rewrite, focusing on performance improvements and leveraging Lean's strengths for formal verification and efficient code. They claim the new version is faster than the original Python implementation, highlighting specific optimizations and the trade-offs involved.
背景速读
- **Tinygrad** 是一个轻量级深度学习框架,主打代码简洁和教育性——它的核心在于用少量代码实现自动微分和 GPU 计算。原版用 Python 和 C 写成,作者是 George Hotz (geohot)。
- **Lean** 是一种交互式定理证明语言,主要用于形式化数学证明(如微软研究院的 Lean 社区)。它并非为机器学习设计,但它的类型系统和高性能编译能力近年被用来做“可证明正确”的系统编程。
- 本文的团队将 Tinygrad 的核心从 Python 重写为 Lean,利用 Lean 的编译器(通过 C 后端)生成更高效的机器码。他们声称在部分基准测试中获得了显著加速——这展示了将形式验证工具引入 ML 基础设施的新趋势。
- 该项目在 GitHub 上开源,引发了对“用证明语言写 ML 框架”的讨论:性能提升是否通用、Lean 是否能替代 CUDA/C++ 的某些角色、形式化方法能否减少训练/推理中的数值 bug。