Source-linked AI summary
Fashionable Modelling with Flux
Michael Innes, Elliot Saba, Keno Fischer, Dhairya Gandhi, Marco Concetto Rudilosso, Neethu Mariya Joy, Tejan Karmali, Avik Pal, Viral Shah
TL;DR
Flux addresses the challenge of building increasingly complex machine-learning models without sacrificing expressiveness or performance. It presents a Julia-based differentiable-programming framework whose simplicity and extensibility are supported by compiler techniques. The paper describes source-level differentiation, accelerator compilation, and automatic batching as mechanisms for combining flexible model code with efficient execution.
Problem
Machine learning models increasingly incorporate control flow, state, data structures, and techniques from multiple computer-science areas, motivating a programming-language perspective.
Method
Flux refines Julia into a differentiable-programming stack combining simple, hackable model code with compiler support for high-performance execution.
Results
Flux supports differentiable programming through source-to-source automatic differentiation, compilation to accelerators, and SPMD-based automatic batching.
Takeaways & Limitations
Flux enables models to combine neural networks with optimization, differential equations, probabilistic programming, and physical simulations within differentiable programs.
Abstract
from arXiv · showhide
Machine learning as a discipline has seen an incredible surge of interest in recent years due in large part to a perfect storm of new theory, superior tooling, renewed interest in its capabilities. We present in this paper a framework named Flux that shows how further refinement of the core ideas of machine learning, built upon the foundation of the Julia programming language, can yield an environment that is simple, easily modifiable, and performant. We detail the fundamental principles of Flux as a framework for differentiable programming, give examples of models that are implemented within Flux to display many of the language and framework-level features that contribute to its ease of use and high productivity, display internal compiler techniques used to enable the acceleration and performance that lies at the heart of Flux, and finally give an overview of the larger ecosystem that Flux fits inside of.
1 Introduction
Flux presents machine learning as differentiable programming expressed through Julia, emphasizing simplicity, hackability, and compiler technology. Its design combines readable, modifiable code with performance across CPUs and accelerators.
- Machine learning models are best viewed as differentiable algorithms, making programming languages central to expressing increasingly complex models.
- Flux uses Julia because its numerical design and extensible compiler support high-performance machine learning on large models and datasets.
- Flux is built around three pillars: simplicity, hackability, and underlying compiler technology.
- 1.1 Simplicity: Flux is approximately a thousand lines of straightforward Julia code, contrasting with typical frameworks written in many hundreds of thousands of lines of C++.
- 1.2 Hackability: Flux treats library code as user code, allowing models to mix mathematical expressions, custom layers, algorithms, and control flow.
- 1.3 Compiler Technology: Flux preserves a dynamic define-by-run interface while using compiler techniques to provide source-to-source differentiation, operator fusion, distributed training, and deployment benefits.
2 Fashionable Modelling
The section illustrates Flux through a discriminative adversarial network for cough classification. The model addresses dataset-label correlation by jointly classifying cough type and dataset provenance.
- Dataset provenance was more strongly associated with recording differences than tuberculosis status was with cough differences within a dataset.
- A Discriminative Adversarial Network adds a provenance classifier that penalizes shared convolutional features useful for identifying the originating dataset.
- The training loop computes classifier and adversarial losses, backpropagates both, and updates the model with an optimizer step.
- The architecture uses a convolutional network for cough type classification and feeds first-layer outputs into a multilayer perceptron for dataset-source classification.
3 Extending Julia’s Compiler
Flux extends Julia’s compiler and differentiable-programming infrastructure to target GPUs, TPUs, automatic batching, and JavaScript while retaining expressive model code. These compiler techniques address accelerator performance and the tradeoff between dynamic control flow and efficient differentiation.
- 3.1 Compiling Julia for GPUs: Julia supports CUDA-style GPU kernels, while type specialization extends them to sparse arrays, complex numbers, and higher-order user-defined operations.Specialization generates appropriate PTX instructions for different element types and array structures.
- 3.1 Compiling Julia for GPUs: Dual numbers enable forward-mode differentiation with stack allocation and fused primal-tangent computation, making derivatives effectively free in memory-bound GPU situations.The approach is particularly valuable for functions containing control flow.
- 3.1 Compiling Julia for GPUs: Element-wise operations can be decomposed into independent scalar computations, allowing sparse Jacobians and derivative calculations to be fused with the original operation.The fused forward-mode results can participate in a reverse-mode sweep to obtain gradients.
- 3.2 Algorithmic Differentiation: Zygote differentiates Julia’s SSA-form syntax directly, supporting control flow, recursion, data structures, and macros while enabling traditional compiler optimization of forward and backward passes.The same infrastructure can support kernel fusion and compilation to accelerators such as TPUs.
- 3.3 Compiling Julia for TPUs: Custom compiler passes identify static Julia sub-segments and compile them into XLA IR for direct TPU execution.Julia’s compiler workflow analyzes, compiles, and runs static sub-segments before returning to dynamic execution when needed.
- 3.3 Compiling Julia for TPUs: VGG19’s forward pass, backward pass, and optimization step can compile into one XLA chunk, avoiding costly host–TPU communication latency.Julia’s type specialization, inference, and constant propagation make these static sub-segments large.
- 3.4 Automatic Batching: SPMD compilation reframes automatic batching as a compiler transformation, allowing code written for individual samples to execute efficiently across batched or SIMD data.This addresses the difficulty of batching variably structured inputs and value-dependent control flow by avoiding manual batching code.
- 3.5 FluxJS: FluxJS partially evaluates Julia model code into JavaScript call graphs, including control flow such as RNNs, and uses tensorflow.js for acceleration.The approach relies on a small set of operator equivalencies between Julia and JavaScript.
4 Conclusion
Flux and Julia provide an environment for high-performance, simple, and hackable machine learning. The paper illustrates readable model expression and compiler techniques that support CPUs and increasingly important accelerators.
- 4 Conclusion: Flux and Julia combine high performance, simplicity, and hackability for machine learning across CPUs and accelerators.The paper presents readable, mathematically recognizable models alongside compiler techniques for performant generated code.