Mentats: Lessons Learned Building a Deep Learning Framework from Scratch in Rust

A Deep Learning Experiment with Zero ML Library Dependencies

In an era where AI frameworks like PyTorch and TensorFlow have virtually monopolized the entry point to deep learning development, one developer chose a harder but far more educational path — building a complete deep learning framework from scratch in Rust, without relying on any external machine learning libraries. The project, named mentats, is now published on crates.io and GitHub.

For the author, this was a double first: their first foray into deep learning and their first deep dive into Rust. They candidly admit that the core motivation behind building mentats was never to compete with mainstream frameworks, but rather to deepen their understanding of the underlying principles through hands-on implementation. Core components like Tensors, Layers, and Optimisers were all built from the ground up, with zero ML dependencies.

Tensors are the most fundamental data structure in deep learning — essentially an abstraction over multi-dimensional arrays that supports efficient numerical operations and automatic differentiation. In frameworks like PyTorch, tensors not only store data but also maintain a Computational Graph that records the dependency of each operation, enabling automatic gradient computation during backpropagation. Implementing tensors from scratch means the developer must manually handle memory layout (row-major vs. column-major), broadcasting mechanics, and the chain rule logic for gradient propagation. Layer implementation involves the mathematical operations of forward propagation and Jacobian matrix computation during backpropagation, while optimizers (such as SGD and Adam) require correctly maintaining state variables like momentum and second-moment estimates, and precisely applying learning rate scheduling strategies at each parameter update step.

This kind of "reinventing the wheel" practice holds special value in the machine learning learning journey. When you have to implement backpropagation, gradient accumulation, and parameter updates by hand, the details hidden behind high-level APIs are exposed in an extremely intuitive way, and your depth of understanding grows accordingly.

From Conditional VAE to Real-World Lessons on Posterior Collapse

The project's most significant achievement so far is the successful training of a Conditional Variational Autoencoder (Conditional VAE) on the MNIST dataset. But the journey was far from smooth — the author encountered one of the most classic and thorny problems in generative modeling: Posterior Collapse.

A Variational Autoencoder (VAE) is a class of generative models based on probabilistic inference, with the core idea of learning the latent distribution of data. A VAE consists of an encoder and a decoder: the encoder maps input data to probability distribution parameters (mean and variance) in a latent space, while the decoder samples from that distribution and reconstructs the data. The training objective is to maximize the Evidence Lower Bound (ELBO), which is equivalent to simultaneously optimizing a reconstruction loss and a KL divergence term. A Conditional VAE (CVAE) extends this by introducing conditional information (such as digit labels in MNIST), enabling the model to generate samples corresponding to specified conditions. MNIST is a classic benchmark dataset containing 70,000 grayscale images of handwritten digits. Due to its moderate size and clear task definition, it has long been regarded as the go-to standard for validating generative and classification models.

What Is Posterior Collapse?

Posterior collapse occurs when the Decoder learns to completely ignore the latent code. Regardless of the input, it outputs an "average-looking" digit. In other words, the model gives up on using information from the encoding space and degenerates into a "lazy" model that only outputs statistical averages. This is one of the most common failure modes in VAE training — the model may appear to converge while actually losing its ability to generate diverse outputs.

Root Cause: The Granularity of Beta Annealing

After thorough investigation, the author identified the main culprit: the beta annealing schedule was computed per epoch rather than per batch.

In VAE training, the weight of the KL divergence term (commonly denoted as β) often needs a "warm-up" process — starting from a small value and gradually increasing, allowing the model to first learn reconstruction before progressively introducing regularization constraints on the latent space. KL divergence (Kullback-Leibler Divergence) measures the distance between the encoder's posterior distribution and the prior distribution (typically a standard normal distribution). In the VAE loss function, the KL term acts as a regularizer, forcing the latent space to maintain structure and continuity. However, if the KL term imposes too strong a constraint early in training, the model tends to let the posterior distribution degenerate directly into the prior — this is precisely the mechanism behind posterior collapse. Beta annealing (also known as KL annealing or cyclical annealing) mitigates this by controlling the growth rate of the coefficient β, setting β to near-zero values at the start of training so the model focuses on learning reconstruction ability, then gradually increasing β to guide the model toward a balance between reconstruction quality and latent space regularization.

The author's original implementation updated the KL weight only once per epoch, rather than continuously updating it based on a global step counter. This made the warm-up schedule far coarser than intended — the granularity of the schedule directly affects how rapidly the constraint strength changes at each training step. The finer the granularity, the smoother the transition, and the less likely the model is to experience abrupt posterior collapse.

After switching the annealing update from "once per epoch" to "continuous per-batch updates," the problem was alleviated. This detail is extremely valuable for beginners — a seemingly insignificant choice in scheduling granularity can directly determine the success or failure of a generative model.

Validation and Training Stability Challenges for Generative Models

During the project's development, the author also posed several highly valuable technical questions to the community — questions that represent common confusions shared by many deep learning practitioners.

A Sanity Check Checklist for VAEs

Before trusting a generative model and committing to large-scale training, what standard sanity checks should you perform? For VAEs, common validation methods include:

• Reconstruction quality check: Observe the model's reconstruction of input samples. If it can't even reconstruct well, there's a fundamental capability issue.
• Latent space interpolation: Perform linear interpolation between the latent encodings of two samples and observe whether the generated results transition smoothly. This verifies whether the latent space has learned meaningful structure. A well-trained VAE should have a continuous latent space — neighboring latent vectors should correspond to visually similar outputs, rather than exhibiting abrupt changes or meaningless noise.
• KL divergence monitoring: Track the KL term and reconstruction loss curves separately. If the KL term rapidly approaches zero, it's often a signal of posterior collapse.
• Random sampling generation: Sample from the prior distribution and decode, checking the diversity and quality of generated samples.

Pitfalls When Transitioning from VAE to GAN

The author's next step is to train a GAN on MNIST and then extend it to a Convolutional GAN. Generative Adversarial Networks (GANs), proposed by Ian Goodfellow in 2014, employ a game-theoretic framework for training generative models. A GAN consists of two adversarial neural networks: a Generator that attempts to produce realistic data samples from random noise, and a Discriminator that tries to distinguish real data from generated data. The two are alternately optimized through a minimax game — the generator aims to maximize the discriminator's error rate, while the discriminator aims to maximize classification accuracy. Ideally, when training converges, samples produced by the generator are indistinguishable from real data, and the discriminator's accuracy approaches 50%. Convolutional GANs (DCGAN) introduce convolutional neural networks into the GAN architecture, replacing fully connected layers with Transposed Convolution for upsampling. This significantly improved image generation quality and laid the foundation for subsequent architectures like StyleGAN and BigGAN.

When transitioning from VAE to GAN, you need to watch out for a series of training stability issues that don't arise with VAEs:

• Mode Collapse: Similar to VAE's posterior collapse but with a different mechanism — the GAN generator may only produce a handful of sample types to fool the discriminator. Specifically, once the generator discovers that certain outputs can consistently deceive the discriminator, it stops exploring other generation modes, resulting in severely lacking output diversity.
• Training imbalance: The capabilities of the discriminator and generator need to maintain a dynamic equilibrium; if one becomes too strong, the other can't learn. Common practical techniques include adjusting the update frequency ratio between the two (e.g., training the discriminator five times for every one generator update), using Label Smoothing, or introducing Spectral Normalization to cap the discriminator's capacity.
• Vanishing gradients: When the discriminator becomes too confident, the gradient signal received by the generator becomes extremely weak. This is an inherent flaw of the original GAN loss function. The subsequently proposed Wasserstein GAN (WGAN) alleviates this by replacing JS divergence with the Wasserstein distance, providing smoother and more informative gradient signals.
• Hyperparameter sensitivity: GANs are far more sensitive to hyperparameters like learning rate and batch size than VAEs. Minor hyperparameter changes can cause training to go from converging to completely diverging — this is the primary reason GANs are considered "hard to tame" in engineering practice.

The insidious nature of these problems is exactly what the author is concerned about — the instabilities of GAN training are hard to anticipate from VAE experience alone.

Advantages and Trade-offs of Implementing Deep Learning in Rust

Choosing Rust as the implementation language is itself a decision worth discussing. Rust was initiated by Mozilla Research in 2010, with version 1.0 released in 2015. Its design goal is to provide memory safety guarantees without sacrificing performance. Compared to Python, Rust offers memory safety guarantees, zero-cost abstractions, and near-C runtime performance — qualities that are quite attractive for low-level, computationally intensive tasks. Through its Ownership, Borrow Checker, and Lifetime mechanisms, Rust eliminates memory safety issues like data races and dangling pointers at compile time rather than runtime, giving programmers memory safety without relying on a garbage collector. Zero-cost Abstraction means that high-level language features incur no additional runtime overhead after compilation — particularly important for deep learning tasks that involve heavy numerical computation.

However, Rust's ML ecosystem is far less mature than Python's. The lack of ready-made automatic differentiation and tensor operation libraries means developers must shoulder significantly more foundational work. Python dominates the ML field not just because of the language's ease of use, but because libraries like NumPy and SciPy, along with frameworks like PyTorch and TensorFlow, have built an extraordinarily rich ecosystem — vast numbers of pretrained models, tutorials, and community resources all revolve around Python. On the positive side, the "inconvenience" of the Rust ecosystem forces developers to understand every computational step rather than treating them as black-box calls. For learning-oriented projects aimed at "deepening understanding," this "inconvenience" actually becomes an advantage.

Notably, Rust is not alone in the deep learning space — native Rust ML frameworks like Candle and Burn are also steadily developing. Candle is a lightweight Rust inference framework from Hugging Face, focused on model deployment scenarios with an emphasis on minimal dependencies and efficient inference. Burn positions itself as a complete deep learning framework supporting multiple compute backends (including CPU, CUDA, and WebGPU), aiming to provide the Rust ecosystem with a development experience similar to PyTorch. While personal from-scratch projects like mentats don't target production readiness, they collectively reflect the community's exploratory interest in safer, higher-performance ML infrastructure.

Conclusion: The Learning Value of Building from Scratch

The mentats project may never become a mainstream framework, but it embodies a learning philosophy worthy of admiration: truly understanding principles through hands-on implementation. From tensor operations to optimizers, from VAE posterior collapse to GAN training stability, every piece of knowledge the author accumulated through this process is more solid and profound than what calling high-level APIs could ever provide.

For anyone looking to deeply understand the internal mechanisms of deep learning, this kind of from-scratch implementation practice is well worth emulating. This methodology has a long tradition in computer science education — just as the best way to understand operating systems is to write a simple kernel, and the best way to understand compilers is to implement a lexer and parser, the best way to understand deep learning is to start from the most basic matrix operations and gradient computations, gradually building up a complete training pipeline. The author's open attitude in proactively seeking community feedback is also a fine embodiment of the open-source spirit.