5 posts tagged with "Machine Learning"

A working notebook on Vision MoE — from V-MoE's hard top-$k$ routing to Soft MoE's continuous slot-based formulation.

A hand-written LSTM with a sparse Mixture-of-Experts FFN between two recurrent layers, trained on Tiny Shakespeare. Anchor: Shazeer et al., Outrageously Large Neural Networks (2017) — the original LSTM-with-MoE paper.

Full runnable notebook: createcentury/jax-flax-optax-lab — 02-lstm-moe.

A hands-on with JAX / Flax / Optax — training a small CNN on CIFAR-10. The full runnable notebook lives in createcentury/jax-flax-optax-lab; this post is the design notes + numbers.

Mamba's (state-spaces/mamba) selective scan is written as a CUDA kernel and doesn't run on Apple Silicon out of the box. I rewrote it in Metal Shading Language (MSL), wired it up to load HuggingFace weights directly, and got end-to-end inference working — that's mamba-metal. This post is the design + results notebook.

Mamba (Gu & Dao 2023) is a sequence model called a Selective State Space Model. It's a candidate replacement for the Transformer's attention, with linear-time compute in context length and a per-token decode cost that is independent of context length. This post walks through the architecture, what "selective" actually means, why the recurrence can still be parallelised, and finally the official CUDA kernel csrc/selective_scan/selective_scan_fwd_kernel.cuh line by line.