3DAlign-DAER

Implementation of 3DAlign-DAER: Dynamic Attention Policy and Efficient Retrieval Strategy for Fine-grained 3D–Text Alignment at Scale (arXiv:2511.13211).

This repository contains training and inference code only. The synthetic data path in daer.data.align3d.SyntheticAlign3D lets the whole pipeline run end-to-end on a laptop / CI with zero external assets, so you can iterate on the method before wiring in a real Align3D-2M manifest.

Method in one paragraph

Given a point cloud and a natural-language caption, DAER encodes each side with a hierarchical encoder (PointNet++-style clustering for 3D, a transformer for text), then fuses them with Hierarchical Attention Fusion (HAF): a stack of cross-attention blocks, one per geometric granularity, whose outputs are combined by a learnable level-weight distribution. The Dynamic Attention Policy calibrates those level-weights during training by running a short MCTS over a small discrete set of per-level deltas; the search is driven by a hybrid reward (Δ recall@1 on the current mini-batch + an entropy / smoothness regulariser). At inference time, retrieval over a large embedding pool uses the Efficient Retrieval Strategy (ERS) — a two-level IVF-style index that probes the top-n_probe clusters and scores only their members, matching exact-KNN recall while being much faster at scale.

Repository layout

3DAlign-DAER/
├── src/daer/
│   ├── data/          # Align3D dataset, augmentations, tokenizer, DataModule
│   ├── models/        # point encoder, text encoder, HAF, Lightning module
│   ├── mcts/          # node, hybrid reward, HAF level-weight calibrator
│   ├── retrieval/     # HierarchicalIndex, ERS, KNN baseline
│   ├── losses/        # bidirectional + hierarchical InfoNCE
│   ├── metrics/       # recall@K, MRR, top-k accuracy
│   └── utils/         # geometry, seeding, logging
├── configs/           # base / pretrain / mcts / retrieval YAMLs
├── scripts/           # train.py, eval.py, infer.py, build_index.py, retrieve.py
├── tests/             # unit + end-to-end (runs on synthetic data)
└── experiments/       # smoke.sh

Quick start

pip install -e .[dev]
pytest -q                                                  # run unit tests
bash experiments/smoke.sh                                  # 1-epoch synthetic smoke
python scripts/train.py --config configs/base.yaml         # full training loop
python scripts/train.py --config configs/mcts.yaml         # + MCTS calibration

Everything above runs on CPU, without downloading any model, by using the synthetic Align3D generator and the built-in bag-of-tokens tokenizer.

Training on real Align3D-2M

1. Produce a JSONL manifest pointing at .npy files of shape (N, 3) or (N, 6) (XYZ + RGB), each line:

   {"points": "shapes/000001.npy", "text": "a red wooden chair", "label": 3, "split": "train"}

2. Switch data.dataset in your YAML to

   data:
     dataset:
       name: align3d
       root: /path/to/shapes
       manifest: /path/to/align3d_manifest.jsonl
       n_points: 1024

3. (Optional) Switch data.tokenizer.kind to hf with a real pretrained text tower:

   data:
     tokenizer: {kind: hf, name: bert-base-uncased, max_len: 64}
   model:
     text_encoder: hf
     text_encoder_kwargs: {name: bert-base-uncased}

4. Run python scripts/train.py --config configs/pretrain.yaml.

Retrieval (ERS vs KNN)

python scripts/build_index.py  --config configs/base.yaml     --output outputs/index.pt
python scripts/retrieve.py     --index  outputs/index.pt      --retrieval-config configs/retrieval.yaml

The retrieve.py script prints a JSON summary with ers_recall@1, knn_recall@1, wall-clock times, and the ERS/KNN speedup.

Configuration schema

data:
  batch_size: 32
  dataset: {name: synthetic | align3d, ...}
  tokenizer: {kind: bag | hf, ...}

model:
  embed_dim: 384
  text_encoder: bag | hf
  point_encoder: {stages: [[k_centroids, k_neighbours, out_dim], ...]}
  haf: {n_heads, dropout}
  mcts: {deltas, max_depth, c_puct, warmup_steps, temperature_*}  # optional
  temperature: 0.07
  haf_weight: 0.5
  lr: 3e-4

trainer:                   # forwarded to pl.Trainer
  max_epochs: 60
  accelerator: auto

Citation

@article{daer2025,
  title  = {3DAlign-DAER: Dynamic Attention Policy and Efficient Retrieval Strategy
            for Fine-grained 3D-Text Alignment at Scale},
  year   = {2025},
  eprint = {arXiv:2511.13211}
}

License

MIT — see LICENSE.