ControlFlow3D

Distilling Multi-View Knowledge into Latent Flow Matching for Point Cloud Upsampling

Official PyTorch implementation of our NeurIPS 2026 paper, including training, evaluation and feature preparation.

Multi-view knowledge at training time. Point-cloud-only upsampling at inference.

Overview | Installation | Reproduction | Training | Evaluation

📰 News

• 🚀 2026-10-24: Code and pretrained weights are released.
• 🗂️ 2026-10-07: ShapeNetPU is released.
• 🎉 2026-09-25: ControlFlow3D is accepted to NeurIPS 2026.

🌟 Overview

Sparse point clouds provide incomplete observations of the underlying surface, making dense geometry ambiguous. ControlFlow3D distills knowledge from pretrained foundation models into a lightweight generation network during training, enabling point-cloud-only inference. The generative process operates in a structured latent space rather than raw coordinates.

The method has three main components:

• Conditional latent flow matching. A frozen Utonia encoder defines the structured latent space, while sparse geometry provides local and global conditions for feature generation.
• Unsupervised shared-prototype distillation. VGGT image tokens and Utonia 3D target tokens are jointly clustered by soft KMeans without pixel-to-point supervision. Prototype consistency aligns the two modalities, while velocity control transfers multi-view knowledge into the base network.
• Mamba-based geometry decoding. A two-stage decoder maps generated latent features to dense coordinates, combining linear-complexity sequence modeling with local geometry modeling.

Mamba-based geometry decoder.

📊 Qualitative Results

4x point cloud upsampling on ShapeNetPU and PU-GAN:

ShapeNetPU

PU-GAN

🗂️ ShapeNetPU Dataset

Download ShapeNetPU and read the dataset documentation on Hugging Face

ShapeNetPU is a ShapeNet-derived dataset for point cloud upsampling, covering 13 object categories.

• 35,762 shapes: 30,395 for training and 5,367 for testing.
• Paired point clouds: 2,048-point sparse inputs and 8,192-point dense ground truth.
• Multi-view training data: eight rendered views per training shape, with view lists and rendering metadata.
• Evaluation geometry: reference meshes for the test split.

The dataset card provides directory layouts, category-level splits, download instructions, file-format details and loading examples.

⚙️ Installation

Reference environment: Linux, Python 3.10, PyTorch 2.3.1 (CUDA 12.1 runtime), CUDA 12.2 toolkit.

git clone https://github.com/I2-Multimedia-Lab/ControlFlow3D.git
cd ControlFlow3D
conda create -n controlflow3d python=3.10 -y
conda activate controlflow3d

bash install.sh

Install CUDA Toolkit first. install.sh installs requirements.txt and builds the CUDA extensions. Run commands from the repository root.

🔁 Reproduction

The workflow is download data → generate features → optionally build cache → train → evaluate. The following commands use the full ShapeNetPU dataset.

Data preparation

Download ShapeNetPU from Hugging Face:

export HF_ENDPOINT=https://huggingface.co
DATASET_ROOT=data/ShapeNetPU
GPU_ID=0

hf download BaldFatTiger/ShapeNetPU --repo-type dataset \
  --local-dir "$DATASET_ROOT"

Raw sample layouts and feature filenames are described in data/README.md.

Feature preparation

Set up separate Utonia and VGGT environments using their official instructions. Place the weights at weights/utonia/utonia.pth and weights/vggt/model.pt; see weight details.

Training uses Utonia point features and VGGT image features. Testing uses sparse Utonia features. Set the encoder Python paths below, then run from controlflow3d:

UTONIA_PYTHON=/path/to/utonia-env/bin/python
VGGT_PYTHON=/path/to/vggt-env/bin/python

python utils/prepare_dataset_features.py \
  --dataset_root "$DATASET_ROOT" --split train \
  --utonia_python "$UTONIA_PYTHON" --vggt_python "$VGGT_PYTHON" \
  --include_gt_features --vggt_feature_reduction mean-pool \
  --batch_size 1 --num_workers 0 --gpu_ids "$GPU_ID" --amp

python utils/prepare_dataset_features.py \
  --dataset_root "$DATASET_ROOT" --split test --encoders utonia \
  --utonia_python "$UTONIA_PYTHON" \
  --batch_size 1 --num_workers 0 --gpu_ids "$GPU_ID" --amp

Memmap cache (optional, recommended): reduces small-file I/O during training. Source-file loading is also supported; see data loading.

python utils/build_memmap_dataset_cache.py \
  --dataset_root "$DATASET_ROOT" --split all \
  --index_mode part --feature_dtype float16

Feature layouts are listed in data/README.md. After updating features, rebuild an existing cache with --overwrite. For pretrained models, use matching projection and preprocessing settings; see feature projections.

🚀 Training

Train on one GPU:

python train.py \
  --dataset_root "$DATASET_ROOT" --gpu_ids "$GPU_ID" \
  --image_feature_source vggt --index_mode part \
  --epochs 300 --batch_size 8 --val_batch_size 8 \
  --num_workers 4 --prefetch_factor 2 --logdir logs

Logs and checkpoints are saved to logs/<run_dir>/:

• last_checkpoint.pth: resume training.
• best_val_checkpoint.pth: best validation CD.

To resume, append --resume "$RESUME" to the training command:

RESUME=$(find logs -name last_checkpoint.pth | sort | tail -n 1)

📏 Evaluation

Select the best-validation checkpoint from the most recent training run:

CHECKPOINT=$(find logs -name best_val_checkpoint.pth | sort | tail -n 1)
python test.py \
  --checkpoint "$CHECKPOINT" --dataset_root "$DATASET_ROOT" \
  --gpu_ids "$GPU_ID" --unit_mode full --inference_index_mode part \
  --full_num_parts 8 --split test --steps 5 \
  --no_copy_mesh_to_gt --save_dir results

For CD/HD, predictions and GT are independently centered and scaled to unit radius. CD sums the two directed mean squared nearest-neighbor distances; HD sums their maxima. P2F measures mean point-to-mesh distance at the original scale for samples with GT meshes.

Results include predicted point clouds, overall/category/sample metrics and run arguments. See python test.py --help for additional options.

📦 Pretrained model

The released ShapeNetPU checkpoint is available from Google Drive. Save it as checkpoint/ControlFlow3D_ShapeNetPU.pth and select it with:

CHECKPOINT=checkpoint/ControlFlow3D_ShapeNetPU.pth

Use the same evaluation command above. Checkpoint format and placement are described in checkpoint/README.md.

📝 Citation

If you find ControlFlow3D or ShapeNetPU useful for your research, please consider citing our work as an acknowledgment:

@misc{liu2026controlflow3d,
  title={ControlFlow3D: Distilling Multi-View Knowledge into Latent Flow Matching for Point Cloud Upsampling},
  author={Yuang Liu and Zhi Zuo and Zhengkai Zhao and Lirui Zhang and Pan Gao and Hao Feng and Zhengzhe Liu},
  year={2026},
  eprint={2610.21190},
  archivePrefix={arXiv},
  primaryClass={cs.CV},
  url={https://arxiv.org/abs/2610.21190},
}

🙏 Acknowledgements

We thank the authors of Utonia, VGGT, ShapeNet, and 3D-R2N2, as well as the point cloud upsampling community, for their open-source contributions.

📜 Terms of Use

Project code is released under the MIT License.

ShapeNetPU is derived from ShapeNet. Please follow the dataset terms of use and the applicable upstream licenses.