Edifice - 92 neural network architectures for Nx/Axon

Your Libraries & OS Mentoring Libraries
, , , , , ,
1

I’ve been working on an open-source framework for training bots for Super Smash Bros. Melee in Elixir (ExPhil). Along the way I got thoroughly distracted by the plethora of architectures to try, benchmarks to run, and optimizations to make.

What came out of that is Edifice – a library of 92 neural network architectures built on Nx and Axon, covering 16 families.

def deps do
[{:edifice, “~> 0.1.1”}]
end

The idea is simple: every architecture follows the same Module.build(opts) pattern and returns an Axon model. There’s also a unified registry so you can do things like:

model = Edifice.build(:mamba, embed_size: 256, hidden_size: 512, num_layers: 4)

Or compare architectures in a loop

for arch ← [:mamba, :retnet, :griffin, :gla] do
model = Edifice.build(arch, embed_size: 256, hidden_size: 512, num_layers: 4)
# … train and evaluate
end

Edifice.list_architectures()

# => [:attention, :bayesian, :capsule, :deep_sets, :densenet, :diffusion, …]

What’s in it:

  • State Space Models: Mamba, Mamba-2 (SSD), S4, S4D, S5, H3, Hyena, BiMamba, GatedSSM, Jamba, Zamba
  • Attention: Multi-Head, GQA, Perceiver, FNet, Linear Transformer, Nystromformer, Performer, RetNet, RWKV-7, GLA, HGRN-2, Griffin/Hawk
  • Recurrent: LSTM, GRU, xLSTM, MinGRU, MinLSTM, DeltaNet, TTT, Titans, Reservoir
  • Vision: ViT, DeiT, Swin, U-Net, ConvNeXt, MLP-Mixer
  • Convolutional: ResNet, DenseNet, TCN, MobileNet, EfficientNet
  • Generative: VAE, VQ-VAE, GAN, DDPM, DDIM, DiT, Latent Diffusion, Consistency Model, Score SDE, Flow Matching, Normalizing Flow
  • Contrastive: SimCLR, BYOL, Barlow Twins, MAE, VICReg
  • Graph: GCN, GAT, GIN, GraphSAGE, Graph Transformer, PNA, SchNet, DeepSets, PointNet
  • Energy/Probabilistic/Memory: EBM, Hopfield, Neural ODE, Bayesian NN, NTM, Memory Networks
  • Meta/Specialized: MoE, Switch MoE, Soft MoE, LoRA, Adapter, Hypernetwork, Capsule, Liquid NN, SNN, ANN2SNN
  • Building Blocks: RMSNorm, SwiGLU, RoPE, ALiBi, PatchEmbed, AdaptiveNorm, CrossAttention, FFN

Pure Elixir, no Python, no ONNX imports. GPU-ready with EXLA. Ships with 17 guides covering ML foundations through architecture theory, plus runnable examples.

Where it’s at:

This is v0.1.1. Every module has been audited against its canonical paper and is backed by 1660+ tests. The architectures I’ve actually used in ExPhil and tried to train on I’m most confident in. I’m sure there are problems I haven’t found yet – the goal is to iron them out and make this something people can reach for when they want to try out a neural net architecture in Elixir.

It’s also just useful as reference. If you want to see how a specific architecture works, the source is there. Point a coding agent at it and ask why you might use a given architecture for whatever you’re working on (I’ve used claude code on various versions of opus to make this jsyk). Or go through a reference implementation and see if you can recreate it: edifice/notebooks/training_mlp.livemd at main · blasphemetheus/edifice · GitHub is an example using the Multi Layer Perceptron (MLP) architecture (a starting point architecture).

PRs, issues, criticisms, praise, and all feedback welcome.

24 Likes
2

Thank you for sharing—I’m trying to convert some models from Python to Axon/Nx, but it isn’t finished yet.

1 Like
3

Would this work with Project M or Rivals of Aether 1/2?

4

So this library is intended to just deal with different neural network architectures. You could implement various architectures to sort of intake gamestate and output controller output (so like what direction a stick is facing, whether buttons are pressed). A system like this can be trained to play a game, given the right training data. For melee, there is a pre-existing method of gathering large amounts of training data (slp files, which act as ways to replay a game after the fact). So if people create a similar replay system for PM or RoA, this library could be used for training models to interface with those games. I don’t know if other problems need to be solved to actually hook up the inputs you’d get from a model to the game itself. Probably.

The gold standard that I’m sort of hoping to replicate in elixir for the exphil library is GitHub - vladfi1/slippi-ai: Imitation learning from slippi replays. And then the philip v2 bot you can just connect to online and play against via ! commands on twitch and Slippi Direct ( Twitch ). I think that is awesome, and if I can just build that capability in Elixir, the exphil library will be a roaring success.

Edifice exists because there’s a huge amount of variety in how to actually implement a neural net, and for elixir, there doesn’t appear to be a library that just implements many different kind of architectures. Or there are ones that sort of give you the building blocks where you can make those architectures (axon, nx, polaris, etc), but the math and theory of how neural nets actually work, and the incredible variation within types of them sort of warrants its own little scaffolding library.

It’s open source, so take a look if you’re trying to use Nx and Axon to do any of these architectures, or think the elixir ecosystem might give you benefits over Python.

2 Likes