AIConfPaper
← Search

Joshua M Susskind

46 accepted papers

2026

Adapting Self-Supervised Representations as a Latent Space for Efficient Generation

Ming Gui, Johannes Schusterbauer, Timy Phan, Felix Krause, Joshua M. Susskind, Miguel Ángel Bautista +1
ICLR 2026poster

We introduce Representation Tokenizer (RepTok), a generative modeling framework that represents an image using a single continuous latent token obtained from self-supervised vision transformers. Building on a pre-trained SSL encoder, we fine-tune only the semantic token embedding and pair it with a…

Cited by 0Source ↗code ↗Scholar ↗
2026

Rethinking JEPA: Compute‑Efficient Video Self-Supervised Learning with Frozen Teachers

Xianhang Li, Chen Huang, Chun-Liang Li, Eran Malach, Joshua M. Susskind, Vimal Thilak +1
ICLR 2026poster

Video Joint Embedding Predictive Architectures (V‑JEPA) learn generalizable off-the-shelf video representations by predicting masked regions in latent space with an exponential moving average (EMA)‑updated teacher. While EMA prevents representation collapse, it complicates scalable model selection a…

Cited by 0Source ↗Scholar ↗
2026

SimpleFold: Folding Proteins is Simpler than You Think

Yuyang Wang, Jiarui Lu, Navdeep Jaitly, Joshua M. Susskind, Miguel Ángel Bautista
ICLR 2026poster

Protein folding models have achieved groundbreaking results typically via a combination of integrating domain knowledge into the architectural blocks and training pipelines. Nonetheless, given the success of generative models across different but related problems, it is natural to question whether t…

Cited by 0Source ↗code ↗Scholar ↗
2026

Text-Conditional JEPA for Learning Semantically Rich Visual Representations

Chen Huang, Xianhang Li, Vimal Thilak, Etai Littwin, Joshua M Susskind
ICML 2026poster

Image-based Joint-Embedding Predictive Architecture (I-JEPA) offers a promising approach to visual self-supervised learning through masked feature prediction. However with the inherent visual uncertainty at masked positions, feature prediction remains challenging and may fail to learn semantic repre…

Cited by 0Source ↗Scholar ↗
2026

To Infinity and Beyond: Tool-Use Unlocks Length Generalization in State Space Models

Eran Malach, Omid Saremi, Sinead Williamson, Arwen Bradley, Aryo Lotfi, Emmanuel Abbe +2
ICLR 2026oral

State Space Models (SSMs) have become the leading alternative to Transformers for sequence modeling tasks. Their primary advantage is efficiency in long-context and long-form generation, enabled by fixed-size memory and linear scaling of computational complexity. We begin this work by showing a simp…

Cited by 0Source ↗Scholar ↗
2025

Denoising Autoregressive Transformers for Scalable Text-to-Image Generation

Jiatao Gu, Yuyang Wang, Yizhe Zhang, Qihang Zhang, Dinghuai Zhang, Navdeep Jaitly +2
ICLR 2025poster

Diffusion models have become the dominant approach for visual generation. They are trained by denoising a Markovian process which gradually adds noise to the input. We argue that the Markovian property limits the model’s ability to fully utilize the generation trajectory, leading to inefficiencies d…

Cited by 0Source ↗PDF ↗Scholar ↗
2025

Flexible Language Modeling in Continuous Space with Transformer-based Autoregressive Flows

Ruixiang ZHANG, Shuangfei Zhai, Jiatao Gu, Yizhe Zhang, Huangjie Zheng, Tianrong Chen +3
NeurIPS 2025poster

Autoregressive models have driven remarkable progress in language modeling. Their foundational reliance on discrete tokens, unidirectional context, and single-pass decoding, while central to their success, also inspires the exploration of a design space that could offer new axes of modeling flexibil…

Cited by 0Source ↗Scholar ↗
2025

INRFlow: Flow Matching for INRs in Ambient Space

Yuyang Wang, Anurag Ranjan, Joshua M. Susskind, Miguel Ángel Bautista
ICML 2025poster

Flow matching models have emerged as a powerful method for generative modeling on domains like images or videos, and even on irregular or unstructured data like 3D point clouds or even protein structures. These models are commonly trained in two stages: first, a data compressor is trained, and in a…

Cited by 0Source ↗PDF ↗Scholar ↗
2025

Mechanisms of Projective Composition of Diffusion Models

Arwen Bradley, Preetum Nakkiran, David Berthelot, James Thornton, Joshua M. Susskind
ICML 2025poster

We study the theoretical foundations of composition in diffusion models, with a particular focus on out-of-distribution extrapolation and length-generalization. Prior work has shown that composing distributions via linear score combination can achieve promising results, including length-generalizati…

Cited by 15Source ↗PDF ↗Scholar ↗
2025

Normalizing Flows are Capable Generative Models

Shuangfei Zhai, Ruixiang ZHANG, Preetum Nakkiran, David Berthelot, Jiatao Gu, Huangjie Zheng +4
ICML 2025oral

Normalizing Flows (NFs) are likelihood-based models for continuous inputs. They have demonstrated promising results on both density estimation and generative modeling tasks, but have received relatively little attention in recent years. In this work, we demonstrate that NFs are more powerful than pr…

Cited by 9Source ↗PDF ↗code ↗Scholar ↗
2025

Parameters vs FLOPs: Scaling Laws for Optimal Sparsity for Mixture-of-Experts Language Models

Samira Abnar, Harshay Shah, Dan Busbridge, Alaaeldin El-Nouby, Joshua M. Susskind, Vimal Thilak
ICML 2025poster

Scaling the capacity of language models has consistently proven to be a reliable approach for improving performance and unlocking new capabilities. Capacity can be primarily defined by two dimensions: the number of model parameters and the compute per example. While scaling typically involves increa…

Cited by 6Source ↗PDF ↗Scholar ↗
2025

Proxy-FDA: Proxy-based Feature Distribution Alignment for Fine-tuning Vision Foundation Models without Forgetting

Chen Huang, Skyler Seto, Hadi Pouransari, Mehrdad Farajtabar, Raviteja Vemulapalli, Fartash Faghri +3
ICML 2025poster

Vision foundation models pre-trained on massive data encode rich representations of real-world concepts, which can be adapted to downstream tasks by fine-tuning. However, fine-tuning foundation models on one task often leads to the issue of *concept forgetting* on other tasks. Recent methods of robu…

Cited by 0Source ↗PDF ↗Scholar ↗
2025

STARFlow: Scaling Latent Normalizing Flows for High-resolution Image Synthesis

Jiatao Gu, Tianrong Chen, David Berthelot, Huangjie Zheng, Yuyang Wang, Ruixiang ZHANG +4
NeurIPS 2025spotlight

We present STARFlow, a scalable generative model based on normalizing flows that achieves strong performance on high-resolution image synthesis. STARFlow's main building block is Transformer Autoregressive Flow (TARFlow), which combines normalizing flows with Autoregressive Transformer architecture…

Cited by 0Source ↗code ↗Scholar ↗
2025

TADA: Improved Diffusion Sampling with Training-free Augmented DynAmics

Tianrong Chen, Huangjie Zheng, David Berthelot, Jiatao Gu, Joshua M. Susskind, Shuangfei Zhai
NeurIPS 2025poster

Diffusion models have demonstrated exceptional capabilities in generating high-fidelity images but typically suffer from inefficient sampling. Many solver designs and noise scheduling strategies have been proposed to dramatically improve sampling speeds. In this paper, we introduce a new sampling m…

Cited by 0Source ↗code ↗Scholar ↗
2025

Target Concrete Score Matching: A Holistic Framework for Discrete Diffusion

Ruixiang ZHANG, Shuangfei Zhai, Yizhe Zhang, James Thornton, Zijing Ou, Joshua M. Susskind +1
ICML 2025poster

Discrete diffusion is a promising framework for modeling and generating discrete data. In this work, we present Target Concrete Score Matching (TCSM), a novel and versatile objective for training and fine-tuning discrete diffusion models. TCSM provides a general framework with broad applicability. I…

Cited by 1Source ↗PDF ↗Scholar ↗
2024

Aggregate-and-Adapt Natural Language Prompts for Downstream Generalization of CLIP

Chen Huang, Skyler Seto, Samira Abnar, David Grangier, Navdeep Jaitly, Joshua M. Susskind
NeurIPS 2024poster

Large pretrained vision-language models like CLIP have shown promising generalization capability, but may struggle in specialized domains (e.g., satellite imagery) or fine-grained classification (e.g., car models) where the visual concepts are unseen or under-represented during pretraining. Prompt l…

Cited by 0Source ↗PDF ↗Scholar ↗
2024

Construction of Paired Knowledge Graph - Text Datasets Informed by Cyclic Evaluation

Ali Mousavi, Xin Zhan, He Bai, Peng Shi, Theodoros Rekatsinas, Benjamin Han +6
COLING 2024main

Datasets that pair Knowledge Graphs (KG) and text together (KG-T) can be used to train forward and reverse neural models that generate text from KG and vice versa. However models trained on datasets where KG and text pairs are not equivalent can suffer from more hallucination and poorer recall. In t…

Cited by 2Source ↗PDF ↗Scholar ↗
2024

Data-free Distillation of Diffusion Models with Bootstrapping

Jiatao Gu, Chen Wang, Shuangfei Zhai, Yizhe Zhang, Lingjie Liu, Joshua M. Susskind
ICML 2024poster

Diffusion models have demonstrated great potential for generating diverse images. However, their performance often suffers from slow generation due to iterative denoising. Knowledge distillation has been recently proposed as a remedy which can reduce the number of inference steps to one or a few, wi…

Cited by 2Source ↗PDF ↗Scholar ↗
2024

Generative Modeling with Phase Stochastic Bridge

Tianrong Chen, Jiatao Gu, Laurent Dinh, Evangelos Theodorou, Joshua M. Susskind, Shuangfei Zhai
ICLR 2024oral

Diffusion models (DMs) represent state-of-the-art generative models for continuous inputs. DMs work by constructing a Stochastic Differential Equation (SDE) in the input space (ie, position space), and using a neural network to reverse it. In this work, we introduce a novel generative modeling frame…

Cited by 14Source ↗PDF ↗Scholar ↗
2024

How JEPA Avoids Noisy Features: The Implicit Bias of Deep Linear Self Distillation Networks

Etai Littwin, Omid Saremi, Madhu Advani, Vimal Thilak, Preetum Nakkiran, Chen Huang +1
NeurIPS 2024poster

Two competing paradigms exist for self-supervised learning of data representations. Joint Embedding Predictive Architectures (JEPAs) is a class of architectures in which semantically similar inputs are encoded into representations that are predictive of each other. A recent successful approach…

Cited by 6Source ↗PDF ↗Scholar ↗
2024

Kaleido Diffusion: Improving Conditional Diffusion Models with Autoregressive Latent Modeling

Jiatao Gu, Ying Shen, Shuangfei Zhai, Yizhe Zhang, Navdeep Jaitly, Joshua M. Susskind
NeurIPS 2024poster

Diffusion models have emerged as a powerful tool for generating high-quality images from textual descriptions. Despite their successes, these models often exhibit limited diversity in the sampled images, particularly when sampling with a high classifier-free guidance weight. To address this issue, w…

Cited by 8Source ↗PDF ↗Scholar ↗
2024

LiDAR: Sensing Linear Probing Performance in Joint Embedding SSL Architectures

Vimal Thilak, Chen Huang, Omid Saremi, Laurent Dinh, Hanlin Goh, Preetum Nakkiran +2
ICLR 2024spotlight

Joint embedding (JE) architectures have emerged as a promising avenue for ac- quiring transferable data representations. A key obstacle to using JE methods, however, is the inherent challenge of evaluating learned representations without access to a downstream task, and an annotated dataset. Without…

Cited by 7Source ↗PDF ↗Scholar ↗
2024

Overcoming the Pitfalls of Vision-Language Model Finetuning for OOD Generalization

Yuhang Zang, Hanlin Goh, Joshua M. Susskind, Chen Huang
ICLR 2024poster

Existing vision-language models exhibit strong generalization on a variety of visual domains and tasks. However, such models mainly perform zero-shot recognition in a closed-set manner, and thus struggle to handle open-domain visual concepts by design. There are recent finetuning methods, such as pr…

Cited by 12Source ↗PDF ↗code ↗Scholar ↗
2024

Pseudo-Generalized Dynamic View Synthesis from a Video

Xiaoming Zhao, R Alex Colburn, Fangchang Ma, Miguel Ángel Bautista, Joshua M. Susskind, Alex Schwing
ICLR 2024poster

Rendering scenes observed in a monocular video from novel viewpoints is a challenging problem. For static scenes the community has studied both scene-specific optimization techniques, which optimize on every test scene, and generalized techniques, which only run a deep net forward pass on a test sce…

Cited by 13Source ↗PDF ↗code ↗Scholar ↗
2024

Scalable Pre-training of Large Autoregressive Image Models

Alaaeldin El-Nouby, Michal Klein, Shuangfei Zhai, Miguel Ángel Bautista, Vaishaal Shankar, Alexander T Toshev +2
ICML 2024poster

This paper introduces AIM, a collection of vision models pre-trained with an autoregressive objective. These models are inspired by their textual counterparts, i.e., Large Language Models (LLMs), and exhibit similar scaling properties. Specifically, we highlight two key findings: (1) the performance…

Cited by 70Source ↗PDF ↗code ↗Scholar ↗
2024

Swallowing the Bitter Pill: Simplified Scalable Conformer Generation

Yuyang Wang, Ahmed A. A. Elhag, Navdeep Jaitly, Joshua M. Susskind, Miguel Ángel Bautista
ICML 2024poster

We present a novel way to predict molecular conformers through a simple formulation that sidesteps many of the heuristics of prior works and achieves state of the art results by using the advantages of scale. By training a diffusion generative model directly on 3D atomic positions without making ass…

Cited by 17Source ↗PDF ↗Scholar ↗
2024

Vanishing Gradients in Reinforcement Finetuning of Language Models

Noam Razin, Hattie Zhou, Omid Saremi, Vimal Thilak, Arwen Bradley, Preetum Nakkiran +2
ICLR 2024poster

Pretrained language models are commonly aligned with human preferences and downstream tasks via reinforcement finetuning (RFT), which refers to maximizing a (possibly learned) reward function using policy gradient algorithms. This work identifies a fundamental optimization obstacle in RFT: we prove…

Cited by 7Source ↗PDF ↗code ↗Scholar ↗
2024

What Algorithms can Transformers Learn? A Study in Length Generalization

Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Joshua M. Susskind +2
ICLR 2024poster

Large language models exhibit surprising emergent generalization properties, yet also struggle on many simple reasoning tasks such as arithmetic and parity. In this work, we focus on length generalization, and we propose a unifying framework to understand when and how Transformers can be expected to…

Cited by 132Source ↗PDF ↗Scholar ↗
2024

When can transformers reason with abstract symbols?

Enric Boix-Adserà, Omid Saremi, Emmanuel Abbe, Samy Bengio, Etai Littwin, Joshua M. Susskind
ICLR 2024poster

We investigate the capabilities of transformer models on relational reasoning tasks. In these tasks, models are trained on a set of strings encoding abstract relations, and are then tested out-of-distribution on data that contains symbols that did not appear in the training dataset. We prove that fo…

Cited by 10Source ↗PDF ↗code ↗Scholar ↗
2023

Diffusion Probabilistic Fields

Peiye Zhuang, Samira Abnar, Jiatao Gu, Alex Schwing, Joshua M. Susskind, Miguel Ángel Bautista
ICLR 2023poster

Diffusion probabilistic models have quickly become a major approach for generative modeling of images, 3D geometry, video and other domains. However, to adapt diffusion generative modeling to these domains the denoising network needs to be carefully designed for each domain independently, oftentimes…

Cited by 28Source ↗PDF ↗Scholar ↗
2023

MAST: Masked Augmentation Subspace Training for Generalizable Self-Supervised Priors

Chen Huang, Hanlin Goh, Jiatao Gu, Joshua M. Susskind
ICLR 2023poster

Recent Self-Supervised Learning (SSL) methods are able to learn feature representations that are invariant to different data augmentations, which can then be transferred to downstream tasks of interest. However, different downstream tasks require different invariances for their best performance, so…

Cited by 11Source ↗PDF ↗Scholar ↗
2023

NerfDiff: Single-image View Synthesis with NeRF-guided Distillation from 3D-aware Diffusion

Jiatao Gu, Alex Trevithick, Kai-En Lin, Joshua M. Susskind, Christian Theobalt, Lingjie Liu +1
ICML 2023poster

Novel view synthesis from a single image requires inferring occluded regions of objects and scenes whilst simultaneously maintaining semantic and physical consistency with the input. Existing approaches condition neural radiance fields (NeRF) on local image features, projecting points to the input i…

Cited by 182Source ↗PDF ↗Scholar ↗
2023

PLANNER: Generating Diversified Paragraph via Latent Language Diffusion Model

Yizhe Zhang, Jiatao Gu, Zhuofeng Wu, Shuangfei Zhai, Joshua M. Susskind, Navdeep Jaitly
NeurIPS 2023poster

Autoregressive models for text sometimes generate repetitive and low-quality output because errors accumulate during the steps of generation. This issue is often attributed to exposure bias -- the difference between how a model is trained, and how it is used during inference. Denoising diffusion mod…

Cited by 20Source ↗PDF ↗code ↗Scholar ↗
2023

Stabilizing Transformer Training by Preventing Attention Entropy Collapse

Shuangfei Zhai, Tatiana Likhomanenko, Etai Littwin, Dan Busbridge, Jason Ramapuram, Yizhe Zhang +2
ICML 2023poster

Training stability is of great importance to Transformers. In this work, we investigate the training dynamics of Transformers by examining the evolution of the attention layers. In particular, we track the attention entropy for each attention head during the course of training, which is a proxy for…

Cited by 71Source ↗PDF ↗code ↗Scholar ↗
2023

Transformers learn through gradual rank increase

Enric Boix-Adserà, Etai Littwin, Emmanuel Abbe, Samy Bengio, Joshua M. Susskind
NeurIPS 2023poster

We identify incremental learning dynamics in transformers, where the difference between trained and initial weights progressively increases in rank. We rigorously prove this occurs under the simplifying assumptions of diagonal weight matrices and small initialization. Our experiments support the the…

Cited by 26Source ↗PDF ↗Scholar ↗
2023

f-DM: A Multi-stage Diffusion Model via Progressive Signal Transformation

Jiatao Gu, Shuangfei Zhai, Yizhe Zhang, Miguel Ángel Bautista, Joshua M. Susskind
ICLR 2023poster

Diffusion models (DMs) have recently emerged as SoTA tools for generative modeling in various domains. Standard DMs can be viewed as an instantiation of hierarchical variational autoencoders (VAEs) where the latent variables are inferred from input-centered Gaussian distributions with fixed scales a…

Cited by 28Source ↗PDF ↗Scholar ↗
2022

Efficient Representation Learning via Adaptive Context Pooling

Chen Huang, Walter Talbott, Navdeep Jaitly, Joshua M Susskind
ICML 2022spotlight

Self-attention mechanisms model long-range context by using pairwise attention between all input tokens. In doing so, they assume a fixed attention granularity defined by the individual tokens (e.g., text characters or image pixels), which may not be optimal for modeling complex dependencies at high…

Cited by 6Source ↗PDF ↗Scholar ↗
2022

GAUDI: A Neural Architect for Immersive 3D Scene Generation

Miguel Ángel Bautista, Pengsheng Guo, Samira Abnar, Walter Talbott, Alexander T Toshev, Zhuoyuan Chen +6
NeurIPS 2022accept

We introduce GAUDI, a generative model capable of capturing the distribution of complex and realistic 3D scenes that can be rendered immersively from a moving camera. We tackle this challenging problem with a scalable yet powerful approach, where we first optimize a latent representation that disent…

Cited by 142Source ↗PDF ↗code ↗Scholar ↗
2022

Learning Representation from Neural Fisher Kernel with Low-rank Approximation

Ruixiang ZHANG, Shuangfei Zhai, Etai Littwin, Joshua M. Susskind
ICLR 2022poster

In this paper, we study the representation of neural networks from the view of kernels. We first define the Neural Fisher Kernel (NFK), which is the Fisher Kernel applied to neural networks. We show that NFK can be computed for both supervised and unsupervised learning models, which can serve as a u…

Cited by 5Source ↗PDF ↗Scholar ↗
2022

Position Prediction as an Effective Pretraining Strategy

Shuangfei Zhai, Navdeep Jaitly, Jason Ramapuram, Dan Busbridge, Tatiana Likhomanenko, Joseph Y Cheng +4
ICML 2022spotlight

Transformers \cite{transformer} have gained increasing popularity in a wide range of applications, including Natural Language Processing (NLP), Computer Vision and Speech Recognition, because of their powerful representational capacity. However, harnessing this representational capacity effectively…

Cited by 23Source ↗PDF ↗code ↗Scholar ↗
2021

Hypersim: A Photorealistic Synthetic Dataset for Holistic Indoor Scene Understanding

Mike Roberts, Jason Ramapuram, Anurag Ranjan, Atulit Kumar, Miguel Angel Bautista, Nathan Paczan +2
ICCV 2021poster

For many fundamental scene understanding tasks, it is difficult or impossible to obtain per-pixel ground truth labels from real images. We address this challenge by introducing Hypersim, a photorealistic synthetic dataset for holistic indoor scene understanding. To create our dataset, we leverage a…

Cited by 395PDF ↗code ↗Scholar ↗
2021

Uncertainty Weighted Actor-Critic for Offline Reinforcement Learning

Yue Wu, Shuangfei Zhai, Nitish Srivastava, Joshua M Susskind, Jian Zhang, Ruslan Salakhutdinov +1
ICML 2021spotlight

Offline Reinforcement Learning promises to learn effective policies from previously-collected, static datasets without the need for exploration. However, existing Q-learning and actor-critic based off-policy RL algorithms fail when bootstrapping from out-of-distribution (OOD) actions or states. We h…

Cited by 224Source ↗PDF ↗code ↗Scholar ↗
2021

Unconstrained Scene Generation With Locally Conditioned Radiance Fields

Terrance DeVries, Miguel Angel Bautista, Nitish Srivastava, Graham W. Taylor, Joshua M. Susskind
ICCV 2021poster

We tackle the challenge of learning a distribution over complex, realistic, indoor scenes. In this paper, we introduce Generative Scene Networks (GSN), which learns to decompose scenes into a collection of many local radiance fields that can be rendered from a free moving camera. Our model can be us…

Cited by 160PDF ↗code ↗Scholar ↗