TopoDiT — Optimize Any Topology 2
TopoDiT is the generative model of Optimize Any Topology 2 (OAT2): a 688.6M-parameter conditional diffusion transformer, trained with flow matching, that generates minimum-compliance structural topologies in the latent space of the frozen OpenTO/NFAE neural-field autoencoder. It is the successor of the latent-diffusion U-Net of Optimize Any Topology (NeurIPS 2025) and takes the same inputs: domain shape, mesh cell size, target volume fraction, boundary conditions and loads — at any resolution and aspect ratio.
Code: the OptimizeAnyTopology2 repository (training, evaluation, GPU FEM).
Results (OpenTO test split, 5,000 problems, zero-shot, no CFG)
| OAT (2025) | TopoDiT / OAT2 | |
|---|---|---|
| failure rate, 1 sample | 40.0% | 21.8% |
| median compliance error, 1 sample | 3.94% | 0.46% |
| failure rate, best of 4 | 25.2% | 11.6% |
| median compliance error, best of 4 | 2.86% | 0.18% |
| failure rate, best of 4 + 10 PGD steps | 14.9% | 5.8% |
Identical FEM and statistics for both rows (CE = (C − C_gt)/C_gt; failure = CE ≥ 100%).
Architecture
- Dense DiT over the 64×64×1 NFAE latent: patch size 4 → 256 tokens, 24 blocks, width 1152, 16 heads; predicts the flow-matching velocity.
- Global conditions (shape, cell size, volume fraction) → one vector concatenated with the timestep embedding into adaLN-Zero modulation. Boundary conditions and loads → attention-pooled spatial tokens (32-cell grid over the domain, every point kept), read by cross-attention in every block. Continuous rectangle position embeddings, QK-normalized attention.
- Condition encoder: 4 layers, width 1152, 12 heads, token width 768.
Training
OpenTO labeled + NITO (894k optimized structures), 50 epochs, 349,200 steps at
effective batch 128, AdamW lr 1e-4 (cosine to 1e-6, 1k warm-up), weight decay 1e-4,
logit-normal timesteps, classifier-free-guidance dropout per condition. Latent
normalization statistics are stored on the model.
Usage
import torch
from datasets import load_dataset
from OAT import NFAE, TopoDiT
from OAT.DataUtils import OpenTO, DiffusionCollator, cached_full_grid_cell
from OAT.Pipelines import FlowMatchPipeline
device = 'cuda'
model = TopoDiT.from_pretrained('OpenTO/TopoDiT').to(device).eval()
nfae = NFAE.from_pretrained('OpenTO/NFAE').to(device).eval()
data = load_dataset('OpenTO/OpenTO', split='test')
ds = OpenTO(data, mode='diffusion', train=False)
batch = DiffusionCollator()([ds[0]]).to(device)
pipe = FlowMatchPipeline(shift=1.0)
with torch.no_grad(), torch.autocast('cuda', torch.bfloat16):
z = pipe.inference(model, batch, num_sampling_steps=20, guidance_scale=1.0)
phi = nfae.decoder(model.denormalize(z.float()))
w, h = data[0]['topology'].size
coord, cell = cached_full_grid_cell(h, w)
density = nfae.renderer(phi, [coord[None].to(device)], [cell[None].to(device)])[0][0, 0]
topology = density.float().cpu().numpy() > 0.5
Recommended sampling: 20 Euler steps, guidance_scale=1.0 (CFG off); 5–10 steps are
equally good zero-shot.
Citation
@inproceedings{nobari2025oat,
title = {Optimize Any Topology: A Foundation Model for Shape- and Resolution-Free Structural Topology Optimization},
author = {Heyrani Nobari, Amin and Regenwetter, Lyle and Picard, Cyril and Han, Ligong and Ahmed, Faez},
booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
year = {2025},
note = {arXiv:2510.23667}
}