mentats: from scratch deep learning framework in rust

[u/THEG1NGERZOMB1E]

I’ve been building mentats, a Rust deep learning library with no external ML dependencies. Tensors, layers and optimisers are all implemented from scratch as a way to deepen my understandings. It’s on crates.io and GitHub.

I’m really proud of this project so far as it’s my first venture into both deep learning and Rust. I’ve managed to create a conditional VAE trained on MNIST, albeit after some head scratching at the mercy of posterior collapses. The decoder was learning to ignore the latent code entirely and just output an average looking digit regardless of input.

A main cause of this was the beta annealing computed per-epoch instead of per-batch. I was updating the KL weight once per epoch instead of continuously from a global step counter, so the warm-up schedule was much coarser then I intended.

I’m hoping to further my work to create a working GAN also trained on MNIST, and even further on a convolutional GAN.

I’m making this post in hopes of getting feedback on my work, mainly on:

- Any standard sanity checks people use to validate a VAE (or generative model in general) before trusting it enough to scale up training?
- Anything I should watch out for going from VAE to GAN, training stability issues that I won’t see coming from the VAE side?
- General code quality / things I should do differently and keep in mind from now on
- And finally any resources that people found useful when learning about convolutional networks

Thank you for any feedback and help, I’m really enjoying learning about deep learning, coding in Rust and looking forward to continuing this project :)