← Search

Yiding Jiang

16 accepted papers

2026

Maximum Likelihood Reinforcement Learning

ICML 2026oral

Maximum likelihood is fundamental to supervised learning but it cannot be directly applied in correctness-based problems with non-differentiable sampling. In these settings, reinforcement learning (RL) is typically used to maximize expected reward. We show that for binary correctness tasks, expected…

Cited by 0SourceScholar
2025

Adaptive Data Optimization: Dynamic Sample Selection with Scaling Laws

ICLR 2025poster

The composition of pretraining data is a key determinant of foundation models' performance, but there is no standard guideline for allocating a limited computational budget across different data sources. Most current approaches either rely on extensive experiments with smaller models or dynamic data…

2025

Learning Parameterized Skills from Demonstrations

NeurIPS 2025poster

We present DEPS, an end-to-end algorithm for discovering parameterized skills from expert demonstrations. Our method learns parameterized skill policies jointly with a meta-policy that selects the appropriate discrete skill and continuous parameters at each timestep. Using a combination of temporal…

Cited by 0SourcecodeScholar
2025

Safety Pretraining: Toward the Next Generation of Safe AI

NeurIPS 2025poster

As large language models (LLMs) are increasingly deployed in high-stakes settings, the risk of generating harmful or toxic content remains a central challenge. Post-hoc alignment methods are brittle: once unsafe patterns are learned during pretraining, they are hard to remove. In this work, we prese…

Cited by 0SourceScholar
2025

Training a Generally Curious Agent

ICML 2025oral

Efficient exploration is essential for intelligent systems interacting with their environment, but existing language models often fall short in scenarios that require strategic information gathering. In this paper, we present **Paprika**, a fine-tuning approach that enables language models to develo…

2024

Understanding prompt engineering may not require rethinking generalization

ICLR 2024poster

Zero-shot learning in prompted vision-language models, the practice of crafting prompts to build classifiers without an explicit training process, has achieved impressive performance in many settings. This success presents a seemingly surprising observation: these methods suffer relatively little fr…

Cited by 11SourcePDFScholar
2023

Neural Functional Transformers

NeurIPS 2023poster

The recent success of neural networks as implicit representation of data has driven growing interest in neural functionals: models that can process other neural networks as input by operating directly over their weight spaces. Nevertheless, constructing expressive and efficient neural functional arc…

2023

On the Importance of Exploration for Generalization in Reinforcement Learning

NeurIPS 2023poster

Existing approaches for improving generalization in deep reinforcement learning (RL) have mostly focused on representation learning, neglecting RL-specific aspects such as exploration. We hypothesize that the agent's exploration strategy plays a key role in its ability to generalize to new environme…

2023

Permutation Equivariant Neural Functionals

NeurIPS 2023poster

This work studies the design of neural networks that can process the weights or gradients of other neural networks, which we refer to as *neural functional networks* (NFNs). Despite a wide range of potential applications, including learned optimization, processing implicit neural representations, ne…

2022

Agreement-on-the-line: Predicting the Performance of Neural Networks under Distribution Shift

NeurIPS 2022accept

Recently, Miller et al. showed that a model's in-distribution (ID) accuracy has a strong linear correlation with its out-of-distribution (OOD) accuracy, on several OOD benchmarks, a phenomenon they dubbed ``accuracy-on-the-line''. While a useful tool for model selection (i.e., the model most likely…

2022

Assessing Generalization of SGD via Disagreement

ICLR 2022spotlight

We empirically show that the test error of deep networks can be estimated by training the same architecture on the same training set but with two different runs of Stochastic Gradient Descent (SGD), and then measuring the disagreement rate between the two networks on unlabeled test data. This builds…

Cited by 145SourcePDFScholar
2022

Learning Options via Compression

NeurIPS 2022accept

Identifying statistical regularities in solutions to some tasks in multi-task reinforcement learning can accelerate the learning of new tasks. Skill learning offers one way of identifying these regularities by decomposing pre-collected experiences into a sequence of skills. A popular approach to ski…

2020

Observational Overfitting in Reinforcement Learning

ICLR 2020poster

A major component of overfitting in model-free reinforcement learning (RL) involves the case where the agent may mistakenly correlate reward with certain spurious features from the observations generated by the Markov Decision Process (MDP). We provide a general framework for analyzing this scenario…

Cited by 171SourceScholar
2019

Predicting the Generalization Gap in Deep Networks with Margin Distributions

ICLR 2019poster

As shown in recent research, deep neural networks can perfectly fit randomly labeled data, but with very poor accuracy on held out data. This phenomenon indicates that loss functions such as cross-entropy are not a reliable indicator of generalization. This leads to the crucial question of how gener…