Tongzhou Wang

Tongzhou Wang

MIT CSAIL

I am a current PhD student at MIT CSAIL working with Antonio Torralba and Phillip Isola on machine learning. My current focuses are geometric structures of learned representations, and enabling efficient, adaptive and general agents via such representations. I am broadly interested in representation learning, reinforcement learning, synthetic training data, dataset distillation, and machine learning in general.

I have spent time at Meta AI working with Yuandong Tian, Amy Zhang, and Simon S. Du. I also collaborate with Alyosha Efros and Jun-Yan Zhu.

Before joining MIT, I worked at Facebook AI Research (now Meta AI) on PyTorch, and studied computer science and statistics at UC Berkeley, where I was fortunate to work with Stuart J. Russell, Ren Ng, and Alyosha Efros.

Click here for my CV.

Open Source Projects

  1. PyTorch core developer (2017 - 2019; team size <10). Data loading, CUDA/CPU kernels, API design, autograd optimization, ML ops, Python binding, etc.
  2. torchreparam developer (2019 - 2020). One of the earliest PyTorch toolkits for re-parametrizing neural networks, e.g., for hyper-nets and meta-learning.
  3. torchqmet developer (2022 - now). PyTorch toolkit for SOTA quasimetric learning.
  4. CycleGAN and pix2pix in PyTorch maintainer (2018 - now). 18.9k stars.

See below for open source code for my researches.

Selected Publications

  1. Improved Representation of Asymmetrical Distances with Interval Quasimetric Embeddings
    [NeurIPS 2022 NeurReps Workshop] [Project Page] [arXiv] [PyTorch Package for Quasimetric Learning]
    Tongzhou Wang, Phillip Isola

    computing-iqe

  2. Denoised MDPs: Learning World Models Better Than The World
    [ICML 2022] [Project Page] [arXiv] [code]
    Tongzhou Wang, Simon S. Du, Antonio Torralba, Phillip Isola, Amy Zhang, Yuandong Tian

  3. On the Learning and Learnability of Quasimetrics
    [ICLR 2022] [Project Page] [arXiv] [OpenReview] [code]
    Tongzhou Wang, Phillip Isola

    quasimetric-function-spaces

  4. Understanding Contrastive Representation Learning through Alignment and Uniformity on the Hypersphere
    [ICML 2020] [Project Page] [arXiv] [code]
    Tongzhou Wang, Phillip Isola

    hypersphere_stl10_scatter_linear_output 
    # bsz : batch size (number of positive pairs)
# d   : latent dim
# x   : Tensor, shape=[bsz, d]
#       latents for one side of positive pairs
# y   : Tensor, shape=[bsz, d]
#       latents for the other side of positive pairs
def align_loss(x, y, alpha=2):
    return (x - y).norm(p=2, dim=1).pow(alpha).mean()
    
def uniform_loss(x, t=2):
    return torch.pdist(x, p=2).pow(2).mul(-t).exp().mean().log()
    PyTorch implementation of the alignment and uniformity losses

  5. Dataset Distillation
    [Project Page] [arXiv] [code] [DD Papers]
    Tongzhou Wang, Jun-Yan Zhu, Antonio Torralba, Alexei A. Efros

    dataset_distillation_fixed_mnist

  6. Meta-Learning MCMC Proposals
    [NeurIPS 2018] [PROBPROG 2018] [ICML 2017 AutoML Workshop Oral] [arXiv]
    Tongzhou Wang, Yi Wu, David A. Moore, Stuart J. Russell

    meta_learning_mcmc_gmm_trace

  7. Learning to Synthesize a 4D RGBD Light Field from a Single Image
    [ICCV 2017] [arXiv]
    Pratul Srinivasan, Tongzhou Wang, Ashwin Sreelal, Ravi Ramamoorthi, Ren Ng light-field-synthesis-pipeline