darsnack

How you interact with parameters.

Lux is similar to Flax (Jax) where the parameters are kept in a separate variable from the model definition, and they are passed in on the forward pass. Notably, this design choice allows Lux to accept parameters built with ComponentArrays.jl which can be especially helpful when working with libraries that expect flat vectors of parameters.

Flux lies somewhere between Jax and PyTorch. Like PyTorch, the parameters are stored as part of the model. Unlike traditional PyTorch, Flux has “functional” conventions, e.g. `g = gradient(loss, model)` vs. `loss.backward()`. Similar to Flax, the model is a tree of parameters.

You would have to add a learning method to tell it how to encode/decode graph data, but the framework is agnostic to the model choice. So any Flux model is supported.

You should check out Makie. Getting it set up can be a bit frustrating if things don’t go right, and there is a small learning curve for using `@lift`, but it is an absolute joy to use once you ramp up.

I use it for my research by default. You can pan, zoom, etc. The subplot/layout system is frankly a lot better than Matlab (and I enjoyed Matlab for plotting!). The best part is that I can insert sliders and drop downs into my plot easily, which means I don’t need to waste time figuring out the best static, 2D plot for my experiment. I just dump all the data into some custom logging struct and use sliders to index into the correct 2D plot (e.g. a heat map changing over time, I just save all the matrices and use the slider to get the heat map at time t).

On the NVidia A100, the standard FP32 performance is 20 TFLOPs, but if you use the tensor cores and all the ML features available then it peaks out at 300+ TFLOPs. Not exactly your question, but a simple reference point.

Now the accelerator in the M1 is only 11 TFLOPs. So it’s definitely not trying to compete as an accelerator for training.

In a more complex example where you actually take a variable, do some operations to it, then reassign it, Pluto.jl encourages you to separate that into multiple cells. The reason is each cell marks a distinct node in the dependency graph. If you prefer to use cells, then the notebook can be smarter about what lines actually need to get re-run and what don't.

A downside to using multiple cells is vertical spacing/visual noise. This is something that the package authors are currently thinking about addressing.

Most of that is not fundamental to Julia or Flux itself. It’s the difference between a monolithic package like TF and source-to-source AD in Julia. The former allows the designers to use their own data structures and external libraries to do optimizations. Source-to-source relies on the underlying IR used by Julia, making optimizations challenging without some compiler assistance. But all of that is in the pipeline with stopgap solutions on the way.

As with most things in Julia, the code developers don’t just want to hack changes that work, but make changes that are flexible, extensible, and can solve many problems at once. So, Flux isn’t ready for prime time yet, but it is definitely worth keeping your eye on it.

I have been using Flux for a year (or more?) and I have never found it to be slower than PyTorch or TF. Granted I am training at most ResNet-20 and mostly smaller models, so maybe there is larger training routines where people have issues. Every single one of these deep learning libraries is mapping to CUDA/BLAS calls under the hood. If you wrote the framework correctly, the performance difference should not be drastically different. And Flux doesn’t have much in terms of overhead. My lab mate uses PyTorch to train the same models as me and his performance is consistently the same or worse.

As for features, I think this is because people coming from TF or PyTorch are used one monolithic package that does everything. That’s intentionally not how Flux or the Julia ecosystem is designed. I’ll admit that there are a lot of preprocessing utility functions that could be better in the larger Julia ML community. But for the most part, the preprocessing required for ML research is available. This is mostly the fault of the community of not having a single document explaining to new users how all the packages work together.

Where the difference between Flux and other ML frameworks is apparent is when you try to do anything other than a vanilla deep learning model. Flux is extensible in a way that the other frameworks are just not. A simple example is with the same lab mate and I trying to recreate a baseline from a paper that involved drawing from a distribution at inference time based on a layer’s output then applying a function to that layer based on the samples drawn. I literally implemented the pseudo code from the paper because in Flux everything is just a function, and chains of models can be looped in a for loop like an array. Dumb pseudo code like statements where you just write for loops are just as fast in Julia. And it was! Meanwhile my friends code came to a grinding halt. He had to resort to numerical approximations for drawing from the distribution because he was forced to only use samplers that “worked well” in PyTorch. This is the disadvantage of a monolithic ML library. I didn’t use “Flux distributions,” I just used the standard distributions package in Julia.

This disadvantage to TF and PyTorch will become even more apparent when you do model-based RL. Flux was designed to be simple and extensible from the start. TF and PyTorch were not.

They even address what you are saying in the article. They mention the looming danger of weaponized AI which is exactly when someone has altered the program to do harm.

That’s the entire point of the article. The AI itself will not develop a need to harm. Users/creators of the AI will program that into it. But AI by itself even with creators of the best intentions has real problems that affect society today. And talking about a hypothetical future has distracted the public narrative from the real problems with AI that we need to solve now.