Getting Started with Julia for Artificial Intelligence and Machine Learning

Julia is a high-performance, dynamic programming language that has been rapidly gaining popularity for numerical computing, data science, and artificial intelligence (AI) / machine learning (ML). It aims to combine the ease of use and expressiveness of Python with the speed of C++.

In this guide, we‘ll explore what makes Julia uniquely suited for AI and ML workloads, walk through getting started with Julia and key packages, and highlight some of the exciting work being done with Julia in the AI/ML community.

Julia‘s Key Advantages for AI and ML

High Performance

One of the main draws of Julia for AI/ML is its excellent performance. Julia code can match or exceed the speed of traditional compiled languages like C/C++, without sacrificing ease of use.

To illustrate, here are some benchmarks comparing the runtime of common ML tasks implemented in Julia vs Python and R:

Sources: [1], [2]

As we can see, Julia consistently outperforms Python and R, often by a significant margin. This is thanks to Julia‘s LLVM-based just-in-time (JIT) compiler, which can generate highly optimized native code on the fly.

Julia also has first-class support for parallelism, multi-threading, and GPU computing, making it easy to scale computation across cores or accelerators. Packages like CUDA.jl provide a high-level interface to NVIDIA GPUs and seamlessly integrate with the rest of the Julia ecosystem.

Automatic Differentiation

Automatic differentiation (AD) is a key building block for many machine learning algorithms, allowing efficient computation of gradients for optimization.

Julia has best-in-class tools for AD, including:

• ForwardDiff.jl: A forward-mode AD package that can differentiate native Julia code with minimal overhead
• Zygote.jl: A source-to-source reverse-mode AD package that supports almost all of Julia‘s language features
• Enzyme.jl: A next-generation AD package that uses LLVM extensions for even faster derivatives

With these tools, it‘s possible to efficiently compute gradients of Julia programs with thousands of parameters. This has enabled a new generation of ML frameworks like Flux.jl that are fully differentiable and can leverage Julia‘s performance.

Probabilistic Programming

Probabilistic programming is an emerging paradigm that allows specifying probabilistic models as regular programs and performs inference automatically. This is a powerful tool for tasks like Bayesian inference, uncertainty quantification and causal reasoning.

Julia has several state-of-the-art probabilistic programming frameworks:

• Turing.jl: A flexible probabilistic programming language with support for Hamiltonian Monte Carlo, variational inference, and more
• Gen.jl: A general-purpose probabilistic programming system with programmable inference
• Soss.jl: A probabilistic programming library that allows composing probabilistic models with neural networks as building blocks

These frameworks leverage Julia‘s metaprogramming capabilities and AD tools to allow intuitive modeling of complex probabilistic processes.

For example, here‘s how we might specify a simple Bayesian linear regression in Turing:

using Turing, StatsPlots

@model function linear_regression(x, y)
    # Priors 
    α ~ Normal(0, 10)
    β ~ Normal(0, 10)
    σ ~ Exponential(1)

    # Model
    ŷ = α .+ β .* x
    y ~ MvNormal(ŷ, σ)
end

# Generate some synthetic data
x = collect(range(-5, 5, length=100))
y = 2x + 3 + randn(100)

# Perform inference
chain = sample(linear_regression(x, y), NUTS(), 1000)

# Plot results
plot(chain)

This fits a Bayesian linear regression model to the data and generates a plot of the posterior parameter distributions with just a few lines of code. The @model macro transforms the Julia function into a probabilistic program that can be sampled from using Markov chain Monte Carlo (MCMC).

Composability

Another advantage of Julia is its emphasis on composability. The language was designed from the ground up to make it easy to create and combine modular, reusable components.

This is especially important in AI/ML, where we often want to experiment with different models, architectures, and algorithms. With Julia, you can seamlessly mix and match functionality from different packages without sacrificing performance.

For example, say we wanted to train a neural network to classify images using a pre-trained ResNet from the Metalhead.jl package and serve the model as an API endpoint. We could easily combine functionality from Flux.jl for training, Metalhead.jl for the vision models, and HTTP.jl for the web server:

using Flux, Metalhead, Images, HTTP

# Load pre-trained ResNet34 model
model = ResNet34()

# Add classification head
model = Chain(model, Dense(512, 10), softmax)

# Load and preprocess data
train_data, test_data = load_dataset()
train_data = [preprocess(x) for x in train_data]

# Train model
loss(x, y) = Flux.crossentropy(model(x), y) 
opt = ADAM()
for epoch in 1:10
    Flux.train!(loss, params(model), train_data, opt)
end

# Define prediction endpoint
prediction(req::HTTP.Request) = begin
    img = preprocess(parse_request(req))
    return model(img) |> argmax
end

# Start web server
HTTP.serve(prediction, "127.0.0.1", 8080)

Here we‘re able to leverage pre-trained models, build the neural net with a familiar Flux API, and serve predictions over HTTP in an idiomatic way. Packages in the Julia ecosystem are designed to work well together.

The Julia ML Ecosystem

The ML ecosystem in Julia has grown significantly in recent years. As of September 2023, the ML & Stats category on JuliaPackages lists over 900 registered packages.

Some key packages to be aware of include:

• Flux.jl: A lightweight ML library that takes Julia‘s expressiveness and applies it to ML models like neural networks and recurrent nets
• MLJ.jl: A standardized ML framework that allows mixing and matching models, data preprocessing, and evaluation strategies
• Knet.jl: A deep learning framework that supports dynamic computational graphs and automatic differentiation using dynamic code generation
• DecisionTree.jl: An implementation of Decision Tree and Random Forest algorithms
• Clustering.jl: A collection of clustering algorithms and utilities
• MultivariateStats.jl: Tools for multivariate statistics and data analysis (e.g. PCA, CCA, LDA)

We can get a sense of how the ML ecosystem has evolved by looking at metrics like the number of ML-related Julia packages over time:

As we can see, the number of ML packages in Julia has seen steady growth, with a notable acceleration in the past few years as key frameworks like Flux and MLJ have matured.

Learning Julia for AI/ML

If you‘re coming to Julia from a Python or R background for data science and ML, here‘s a recommended learning path:

1. Start with the Julia Data Science tutorials to get familiar with Julia‘s syntax, package manager, and key libraries like DataFrames and Plots
2. Work through the Flux Model Zoo examples to see how neural networks are built and trained in Julia
3. Explore the MLJ Tutorials for an overview of the MLJ framework and its functionality
4. Check out the ThinkJulia book for a more comprehensive introduction to Julia programming, with data science applications
5. Dive into the JuliaAcademy courses for in-depth, interactive lessons on topics like Decision Trees or Optimization

Real-world Applications

Julia is increasingly being used for cutting-edge AI/ML research and applications. Notable examples include:

• Self-driving cars: Cruise uses Julia to train and validate the ML models powering their autonomous vehicles
• Robotics: Researchers at MIT, Stanford and other universities are using Julia for robot dynamics, control and reinforcement learning
• Drug discovery: Pumas AI built their pharmaceutical modeling and simulation platform on Julia to achieve GPU-accelerated speedups over legacy R code
• Aerospace: NASA and MIT‘s Lincoln Lab use Julia to optimize aircraft trajectories and design satellite constellations
• Climate modeling: The ClimaCore framework uses Julia‘s GPU and ML capabilities to accelerate climate model development

We have collaborators all over the world running large scale computations in Julia across industries. It‘s being used to fly drones, design antennas, optimize power systems, identify new pharmaceutical drugs, guide telescopes, and much more. The combination of ease of use and raw speed makes it an ideal language for AI and ML.

• Dr. Alan Edelman, Co-creator of Julia Computing

Conclusion

Julia is a powerful and expressive language that is well-suited for the demands of modern AI and machine learning workloads. Its unique combination of performance, ease of use, and state-of-the-art tools for AD and probabilistic programming make it a compelling choice for researchers and practitioners.

While the Julia ML ecosystem is still young compared to Python‘s, it is growing rapidly and already provides a robust set of tools for most common ML tasks. And with increasing industry and academic adoption, the future looks bright for Julia as a language for AI/ML.

To get started with Julia for machine learning, begin with the official Julia documentation, then explore the Julia Data Science site and JuliaML GitHub org for tutorials and examples. Join the Julia Discourse to engage with the community and get help as you learn.

References

1. Avik Pal and Viral B. Shah. 2019. Julia for Machine Learning. arXiv:1909.03154. https://arxiv.org/abs/1909.03154
2. Zhang D and Chen G. 2020. Comparison of the speed performance of programming languages in artificial intelligence algorithms. In Journal of Physics: Conference Series 1684:012116. https://doi.org/10.1088/1742-6596/1684/1/012116