I strongly believe software frameworks cannot be designed in the abstract.
UchenML came out of my experiments integrating machine learning into edge
applications, and out of what I want to build next — distributed, heterogeneous
intelligence.
The first demo, dots, puts a trained neural network in your
browser tab as your opponent. No inference server, no API call — a 418 KB
WebAssembly binary and a weights file, served as static assets. (It still
carries introspection hooks for a future demo; stripped, it gets smaller.)
Behind it: 2,760,322 parameters packed to fp16, and a 128-channel residual trunk
running eleven 3×3 convolutions across the 24×24 board for every position it
evaluates — up to 128 positions per move, all in the browser tab. That model was
trained with UchenML and deployed with UchenML. It is almost certainly bigger
than it needs to be — I spent the training budget on making it stronger, not on
making it smaller.
I am not currently able to publish the source code. There is a
publicly available snapshot, about a year
old. I plan to publish the current codebase, but I cannot commit to a timeline
yet.
I chose C++ for ease of integration and portability. With Python, the model
lives in a separate world, and reaching it from the application that needs it
means writing and maintaining glue. I want the model inside one software
development workflow: same repository, same build, same review, same tests.
The longer-term bet is heterogeneous, distributed intelligence: not one large
model behind an API, but many small specialized ones running where the data
already is — in the browser tab, inside the app, or in a cloud service. That
only works if a model is cheap to embed anywhere, which makes the deployment
target the design constraint rather than an afterthought.
UchenML is C++20, built with Bazel, and tested on Visual C++, GCC, and Clang. CI
runs the suite on Linux, macOS, Windows, and under Emscripten for WebAssembly.
The compute backend is hand-written portable SIMD on Google's Highway, so a
single kernel source covers AVX-512, NEON, and WASM SIMD with no
architecture-specific branch in the library. Small, fast, and easy to drop into
an existing C++ project is the whole design brief.
The model is a variable — I usually make it constexpr
— declared in a header.
Inference, training, and everything else are templated on that variable, so the
same information is never stated twice. You never describe a layer's input type
when the previous layer already declared its output; the compiler works that
out, along with parameter counts and scratch buffer sizes. There are extra
facilities for introspection and for advanced use cases, but the core is a
simple composition of layers.
The framework is split into clearly defined modules, and you only pay for what
you use — convolution, RNN, and attention are separate targets, and depending on
one does not drag in the others. Training mirrors that structure in a parallel
tree, and deliberately stays out of the runtime one.
Models compose with |
:
#include "uchen/layers.h"
#include "uchen/linear.h"
constexpr auto model = uchen::layers::FloatModel<1>
| uchen::layers::Linear<2>
| uchen::layers::Relu
| uchen::layers::Linear<1>;
// (1×2 weights + 2 biases) + (2×1 weights + 1 bias)
static_assert(model.all_parameters_count() == 7);
Linear<2>
declares only its output width; the input width is deduced from
whatever it is piped onto. That is why the same descriptor composes anywhere in
a stack without restating shapes — and why getting it wrong is a compile error.
Because layers are values, a block of architecture is a named constant. This is
the residual block the dots network is built from — two padded 3×3 convolutions
on one branch, identity on the other, summed and activated — with C
channels
over an S
×S
grid:
template <size_t C, size_t S>
inline constexpr auto kResBlock =
uchen::layers::Fork<2>
| uchen::layers::Parallel(uchen::layers::Conv2d<C, 3, 3, 1, 1>
| uchen::layers::Relu
| uchen::layers::Conv2d<C, 3, 3, 1, 1>,
uchen::Layer<>())
| uchen::layers::Join
| uchen::layers::Reduce<uchen::PlusOp, 2>
| uchen::layers::Relu
| uchen::layers::Reshape<uchen::convolution::ConvolutionInput<C, S, S>>;
kResBlock<64, 16>
and kResBlock<96, 24>
are different types with different
parameter counts, both resolved at compile time. A trunk is a constexpr
pipeline of these; multi-head outputs come from Fork
and Parallel
the same
way, with each head's loss supplied separately at training time.
The next post will dig into real-world model definitions and the more advanced
composition they need.
The model variable is a functor, so evaluating it is a call, and a
ModelParameters
object supplies the weights:
uchen::ModelParameters parameters(&model, weights); // std::span<const float>
auto output = model(input, parameters);
An optional third argument preallocates the scratch space. Its size is another
compile-time property of the model, and passing it is what makes the forward
pass allocation-free:
// The context is usually too large for the stack, so it goes on the heap. It
// can equally come from a memory-mapped region, an arena, or anywhere else.
auto context = std::make_unique<uchen::ModelContext<Model>>(parameters);
auto output = model(input, parameters, *context);
One context can be reused across calls. Reusing it is not thread-safe, but the
model is a pure function of its input and parameters, so a context per thread
runs in parallel — and the parameters, read-only throughout, are shared between
them.
The parameter count is a property of the type too, so the weight buffer is a
fixed-size array sized by the compiler — no allocation, no length field to get
wrong:
static constexpr size_t kParamCount = uchen::ModelParameters<Model>::P;
std::array<float, kParamCount> weights{};
The on-disk format is that same flat array of floats, so weights can be memory
mapped rather than parsed. In the browser build the array is filled straight
from the ArrayBuffer
the page fetched: the demo ships its weights fp16-packed
— 5.5 MB on the wire, widened to floats on load, since WASM offers no fp16 SIMD
— and binds them with one span
. There is no model , no operator
registry, no interpreter warm-up. Once the bytes have arrived, the network is
ready.
Training uses the same model definition, with uchen/training/
added. Fitting a
piecewise-linear function, end to end:
#include "uchen/training/training.h"
uchen::training::TrainingData<uchen::Vector<float, 1>,
uchen::Vector<float, 1>> data(samples.begin(),
samples.end());
uchen::training::Training training(&model,
{&model, kInitialWeights},
uchen::training::SgdOptimizer<Model>{});
while (training.Loss(data) > 1e-4f) {
training = training.Generation(data, 0.001f).next;
}
uchen::ModelParameters parameters = training.parameters();
TrainingData
is a collection of input-output pairs. It also handles splitting
a dataset into batches and reserving validation records.
Generation
does not mutate anything — it returns the next training state,
which makes a step cheap to inspect, discard, or checkpoint, and makes the loop
read as a fold over generations. Optimizers are pluggable: SgdOptimizer
(stochastic gradient descent), momentum SGD, and AdamOptimizer
come out of the
box, alongside Kaiming-He initialization, per-layer gradient statistics, and
multi-threaded batch gradients.
The loss function is deduced from the model's output type, and the framework
provides a few common ones. Pass your own to the Training
constructor instead
and its gradient is computed automatically through the chain.
What comes out is a flat float array — and that is the entire deployment
format, the same bytes the browser build maps into its std::array
. Training
and serving cannot disagree, because there is nothing for them to disagree
about.
Training
is multi-threaded, but it does not by itself scale to real-world
runs. The next post will discuss the primitives it is built from, and how I used
them to run a larger, more reliable training pipeline with Bazel as the executor.
Adding a layer means writing a type the composition machinery recognizes. The
runtime contract is small:
class MyLayer {
public:
using input_t = uchen::Vector<float, 8>;
using output_t = uchen::Vector<float, 4>;
static constexpr size_t parameter_count = 36;
using scratch_area_t = std::array<float, 4>;
output_t operator()(const input_t& input,
const uchen::Parameters& parameters,
memory::LayerContext<scratch_area_t>* context) const;
};
scratch_area_t
comes from the model context, and the framework controls its
lifetime. The layer uses it for its output and for any intermediate buffers it
needs.
operator()
is reached through a dispatcher that asks the compiler which of
those signatures your layer actually defines, so a layer takes only what it
needs — the input alone, the input and its parameters, or those plus a scratch
context. Composite layers own no storage; they slice the model's single flat
parameter blob at compile-time offsets, which is why a Fork
of four heads
still costs exactly one allocation-free array.
A pipeable descriptor turns the type into something that composes:
inline constexpr uchen::Layer<MyLayerDesc> MyLayer;
The descriptor's job is to map "the model so far" onto a concrete layer type,
binding the input shape the caller never had to write down. It hooks in through
an ADL-found StackLayer(model, desc)
. It is also where optimizations hook in —
kernel fusion, or eliding a layer that is a no-op at the given input shape.
Gradients are a separate, optional surface. A layer becomes trainable when an
overload of ComputeGradients
for it is visible:
auto ComputeGradients(const MyLayer& layer, const MyLayer::input_t& input,
const Grad& output_gradient,
const uchen::Parameters& parameters,
std::span<float, MyLayer::parameter_count> gradients,
const void* scratch);
Every layer in the framework — convolution, RMSNorm, pooling, RNN, softmax —
arrived through exactly this path. There is no privileged built-in set.
UchenML is still very much a work in progress, but it has reached the point
where I can hand agentic tools a model and let them build and train it — which
frees me to spend more time on the application side. The feature list is long.
Among the things I want to add:
I plan to post every few weeks. Next up is a deep dive into the dots project and
how UchenML held up in the real world, followed by articles on the optimization
techniques inside the framework, kernel fusion among them. More demos are in the
works too.