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}
}