README.md

---
license: mit
library_name: pytorch
tags:
- topology-optimization
- structural-optimization
- flow-matching
- diffusion-transformer
- pytorch_model_hub_mixin
- model_hub_mixin
datasets:
- OpenTO/OpenTO
pipeline_tag: image-to-image
---

# 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](https://huggingface.co/OpenTO/NFAE) neural-field autoencoder. It is the
successor of the latent-diffusion U-Net of
[Optimize Any Topology](https://arxiv.org/abs/2510.23667) (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

```python
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

```bibtex
@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}
}
```