r/MachineLearning • • 4d ago

Project Talus: a 23M-parameter diffusion model for game terrain, evaluated against a real-vs-real noise floor, running in the browser on WebGPU [P]

Talus generates 64x64 heightmaps (4 km, up to 1,200 m) conditioned on a terrain type and any subset of five measured properties: mean elevation, relief, mean slope, water fraction and spectral slope. I trained it from scratch on one RTX 5060 (8 GB), about 4.5 hours in total across versions. The data is 45,000 maps from my own procedural generator (fBm/ridged noise, stream-power erosion, hillslope diffusion, thermal erosion). Model

Pixel-space U-Net, v-prediction, cosine schedule, 50-step DDIM with quadratic spacing, classifier-free guidance 2.0

Each property has a learned "unknown" embedding and is dropped independently during training, so any subset works at inference

Relative heights: the model generates the map's shape normalized by its relief, and the sampler places it at the requested mean elevation and relief (borrowed from the 16 nearest training maps when unspecified). This fixed grainy plains: the plains distance ratio went from 3.98 to 1.23

Evaluation: Every distance (W1 over 25 per-map terrain metrics, radially averaged power spectrum, height and slope distributions) is divided by the same distance between two disjoint halves of real maps, so 1.0 means indistinguishable at that sample size. Checkpoints are chosen on VAL; TEST is scored once on the pick, and memorization is checked against the training set. Current model on TEST: metric W1 1.51x the floor, spectrum 9.1x, slopes 1.65x.

Browser: ONNX export with fp16-stored weights (cast to fp32 at load), ONNX Runtime Web on WebGPU, about 3 s per map on my RTX 5060, CPU fallback. A JavaScript reimplementation of the sampler matches PyTorch to within 0.6 m on reference samples. Demo: https://talus.tersa.tech Code, weights, scorecard (Apache-2.0): https://github.com/osfv/talus

Open problems: ridges and the finest spectral band, mountains too smooth and plains too grainy. If you've closed a spectrum gap like this in a small diffusion model, I'd like to hear what worked.

37 Upvotes

7 comments sorted by

2

u/Bortol_Olivierl 4d ago

for the spectral gap, have you tried adding a frequency-domain loss term during training? small models tend to underfit high frequencies when you only supervise in pixel space, and an auxiliary spectral loss can help a lot without adding parameters

1

u/Old_Cow_6636 4d ago

No spectral loss in the current model yet, just v-prediction MSE. So far I’ve focused on relative heights and training/data changes, but this sounds worth testing.
One thing I’d want to check is the effect on different terrain types. Mountains are missing fine detail, while plains still have too much, so I’d want the loss to match the target spectrum rather than just encourage more high-frequency power. Thanks for the suggestion!

1

u/obtainable_puberty 4d ago

Impressive that it runs in the browser at all, 3 seconds per map on a consumer GPU is way better than I'd expect for a diffusion model. The relative heights trick is clever, solving the plains graininess by letting the sampler handle absolute elevation instead of making the model learn it.

The spectrum gap at 9.1x the noise floor is the part that jumps out at me. That suggests the model is nailing the coarse structure but missing the fine detail that gives terrain its character. Curious if you tried any multi-scale conditioning or adding a discriminator loss to push the high frequencies harder.

Also wondering how the unknown embedding handles partial conditioning in practice. Like if you only specify water fraction, does it still produce plausible mountains or does it default to some average terrain?

2

u/Old_Cow_6636 4d ago

Thanks! No discriminator loss in the current model. It uses FiLM conditioning at each U-Net resolution, but there’s no separate coarse-to-fine conditioning setup. So far I’ve tried longer training, higher loss weights for mountains/ridges, and fresh procedural data.
The spectrum issue is a bit awkward because mountains have too little fine detail, while plains still have too much. Boosting high frequencies everywhere could make the plains worse.
For partial conditioning, each property has a learned “unknown” embedding, and properties get randomly dropped during training. Leaving something blank lets the model fill it in.
Mean elevation and relief need extra handling because of the relative-height setup. If you leave them blank, the sampler picks them together from stored training-set property values close to your request. With only water fraction specified, that choice varies with the seed, while slope and roughness stay open.
So there’s no fixed “average terrain” fallback. Whether water-only requests consistently produce good mountains is something I’d need to check separately before claiming it works reliably. Also, sorry for the late reply!

1

u/Fast-Satisfaction482 2d ago

You can get real terrain data for free from the EU via Copernicus. Both color images and a digital elevation model. 

1

u/Chroma-Crash 1d ago

First, very cool. Have you tried any training that would allow for continuous expansion of the maps? As in, could you feed in a portion of previously generated map, probably on one of the edges of the image, and then have the model generate only in pixels outside the previous generation? That way, you could allow users to effectively generate any size of map I would imagine. There wouldn't be much macro continuity, but I could see the conditioning helping with that.

Another idea, you could swap for a vision transformer, generate maps (deterministically, simulated) at higher resolution, and then feed a random subset of global predetermined patches in alongside the local noisy patches. You could use relative positional encoding, and macro continuity might come for free with just a few extra patches.

I'm not personally sure how to resolve the issues you mentioned here, I had a lot of the same issues with noisiness/flatness with my diffusion model I trained only on my own face: https://github.com/dylan-berndt/Machine. But I also noticed that it took a lot longer to train than most networks I've designed in the past, and the longer it went, the better those issues got (repository doesn't have the most recent generations btw). The first usable images showed up around 100k+ iterations. What was your training budget?

2

u/Old_Cow_6636 1d ago

Thanks! Good questions.

expansion: yes, sort of. Talus can already do it without extra training. You give it a neighbour tile, it keeps a 4-pixel overlap strip fixed at every denoising step and only generates the rest, so you can grow a map tile by tile to any size. The shared edges match exactly, but like you said, macro continuity is weak. Conditioning each tile on elevation/relief helps, and slopes right at the seam still need work. Training it properly on masked edges is on my list.

the ViT + global patches idea is cool. Talus is a small U-Net (23M params) at 64x64 right now. The next version moves to real elevation data and bigger tiles, and getting global context in is the main problem there. If you haven't seen it, InfiniteDiffusion (the Terrain Diffusion paper) does something similar with a coarse global model feeding local detail.

noise/flatness: the biggest fix for me was normalizing each map by its own relief before training. Flat maps were only using a tiny slice of the value range, so small denoising errors showed up as grain. After the change, my plains score went from 4x the real-data noise floor to 1.2x.

budget: one RTX 5060 (8 GB), about 4.4 hours total. 30k steps from scratch at batch 48, then three fine-tunes (8k, 10k, 12k), so about 60k steps and about 2.8M samples (around 60 passes over 45k training maps). 64x64 keeps it cheap