How a Browser Trains a Neural Network: Tensors as Textures, Ops as Shaders
You can open a tab, click a button, and watch an LSTM's loss curve go down. No server, no Python environment, no CUDA. How does that work? The short answer: the browser borrows your graphics card through WebGL, and a fe
You can open a tab, click a button, and watch an LSTM's loss curve go down. No server, no Python environment, no CUDA. How does that work?
The short answer: the browser borrows your graphics card through WebGL, and a few clever mappings turn a graphics API into a math engine.
Why a GPU at all?
Training a neural network is mostly matrix multiplication: the same simple operation repeated over thousands of numbers. A CPU does a few things at a time very fast. A GPU does thousands of simple things at the same time. That is a good match.
The catch is that WebGL was designed for drawing triangles, not for multiplying matrices.
The trick: tensors become textures, ops become shaders
Libraries like TensorFlow.js use two mappings:
| Machine learning | WebGL |
|---|---|
| Tensor (an array of numbers) | Texture (an image stored in GPU memory) |
| Operation (matmul, add, ReLU) | Fragment shader (a small program that runs once per pixel) |
Think of a texture as a grid of pixels, where each pixel's color channels hold numbers instead of colors. To "compute" something, the library draws a rectangle that covers the output texture. The GPU runs the fragment shader once for every output pixel, in parallel, and each run computes one element of the result.
JavaScript -> WebGL API -> GPU driver -> graphics hardware
Nothing is being drawn on screen. The "image" is the answer.
What one training step looks like
- Forward pass. Input tensors are uploaded as textures. Each layer is a chain of shader programs that read one texture and write another.
- Loss. A final shader reduces the prediction error to a number.
- Backward pass. Gradients are just more operations, so they also run as shaders.
- Update. The optimizer (SGD, Adam) updates the weight textures, again on the GPU.
The important detail: data stays on the GPU between steps. Only small things, like the loss value you want to print, come back to JavaScript.
A minimal example
import * as tf from '@tensorflow/tfjs';
await tf.setBackend('webgl');
await tf.ready();
const model = tf.sequential();
model.add(tf.layers.lstm({ units: 8, inputShape: [5, 1] }));
model.add(tf.layers.dense({ units: 1 }));
model.compile({ optimizer: tf.train.adam(0.01), loss: 'meanSquaredError' });
await model.fit(xs, ys, {
epochs: 5,
batchSize: 32,
callbacks: {
onEpochEnd: (epoch, logs) => console.log(epoch, logs.loss),
},
});
Here xs and ys are tensors of shape [samples, 5, 1] and [samples, 1]. Everything between fit and the callback runs through the mappings above.
Where it hurts
WebGL gets the job done, but it has real limits.
-
Reading data back is slow. Calling something like
dataSync()forces the CPU to wait for the GPU. Do it once per epoch, not once per step. - First run is slower. Shaders have to be compiled the first time each operation is used.
-
Memory is not garbage collected. Textures live on the GPU. Wrap temporary work in
tf.tidy()and calldispose()on tensors you no longer need, or memory will creep up. - Texture size and precision limits. Large tensors are split across textures, and some devices, especially phones, may use lower precision than you expect.
- No general-purpose compute. WebGL has no compute shaders, so every operation is squeezed into a rendering pass. Newer APIs such as WebGPU are designed to fix this, but browser support still varies, so check before relying on it.
- Background tabs get throttled. Browsers slow down timers and pause animation frames in hidden tabs, so a training loop can stall when you switch away. Keeping the tab in the foreground is the simple fix.
Practically, this means small to medium models: LSTMs with a few hundred units, small Transformers, time-series experiments. Not billion-parameter models.
Why bother training in the browser?
- Privacy. The data never leaves the user's device.
- Zero server cost. The user's GPU does the work.
- Instant demos. A link is all anyone needs.
The trade-off is scale. For learning, prototyping, and visualizing how training behaves, it is hard to beat.
Try it yourself
I built a small tool around this idea: you can train LSTM and Transformer sequence models on time-series data entirely in your browser, with live loss curves, at ai-spinach.xyz. The write-up above is part of its technical docs.
If you have hit different limits with WebGL training, or you are experimenting with WebGPU, I would like to hear about it in the comments.
Originally published by Dev.to WebDev. Aggregated on AIWithGhost for educational purposes — full credit and traffic to the original publisher.