AIS-Bench/AIS-Bench huggingface dataset:
Training a Generally Curious Agent
Fahim Tajwar * 1 Yiding Jiang * 1 Abitha Thankaraj 1 Sumaita Sadia Rahman 2 J Zico Kolter 1 Jeff Schneider 1
Russ Salakhutdinov 1
Abstract
Efficient exploration is essential for intelligent
systems interacting with their environment, but
existing language models often fall short in sce-
narios that require strategic information gathering.
In this paper, we present PAPRIKA, a fine-tuning
approach that enables language models to develop
general decision-making capabilities that are not
confined to particular environments. By training
on synthetic interaction data from different tasks
that require diverse strategies, PAPRIKA teaches
models to explore and adapt their behavior on
a new task based on environment feedback in-
context without more gradient updates. Experi-
mental results show that models fine-tuned with
PAPRIKA can effectively transfer their learned
decision-making capabilities to entirely unseen
tasks without additional training. Unlike tradi-
tional training, our approach’s primary bottleneck
lies in sampling useful interaction data instead
of model updates. To improve sample efficiency,
we propose a curriculum learning strategy that
prioritizes sampling trajectories from tasks with
high learning potential. These results suggest
a promising path towards AI systems that can
autonomously solve novel sequential decision-
making problems that require interactions with
the external world.
1. Introduction
Large language models (LLMs) are considered to be a
promising foundation for autonomous agents, systems capa-
ble of achieving goals independently with minimal human
supervision or intervention. A crucial requirement for such
systems is the ability to interact effectively with external en-
vironments and gather the information necessary to achieve
their objectives. This capability can be formalized as solving
*Equal contribution
1CMU 2North Carolina State University.
Correspondence to: Fahim Tajwar <ftajwar@cs.cmu.edu>.
Proce...
Statistical Test for Feature Selection Pipelines by Selective Inference
Tomohiro Shiraishi * 1 2 Tatsuya Matsukawa * 1 Shuichi Nishino 1 2 Ichiro Takeuchi 1 2
Abstract
A data analysis pipeline is a structured sequence
of steps that transforms raw data into meaning-
ful insights by integrating various analysis algo-
rithms. In this paper, we propose a novel statis-
tical test to assess the significance of data analy-
sis pipelines. Our approach enables the system-
atic development of valid statistical tests applica-
ble to any feature selection pipeline composed of
predefined components. We develop this frame-
work based on selective inference, a statistical
technique that has recently gained attention for
data-driven hypotheses. As a proof of concept,
we focus on feature selection pipelines for linear
models, composed of three missing value imputa-
tion algorithms, three outlier detection algorithms,
and three feature selection algorithms. We theoret-
ically prove that our statistical test can control the
probability of false positive feature selection at
any desired level, and demonstrate its validity and
effectiveness through experiments on synthetic
and real data. Additionally, we present an im-
plementation framework that facilitates testing
across any configuration of these feature selection
pipelines without extra implementation costs.
1. Introduction
In practical data-driven decision-making tasks, integrating
various types of data analysis steps is crucial for address-
ing diverse challenges. For instance, in genetic research
aimed at identifying genes linked to a specific disease, the
process often begins with preprocessing tasks such as filling
in missing values and detecting outliers. This is followed
by screening for potentially related genes using simple de-
scriptive statistics and then applying more complex machine
learning-based feature selection algorithms. Such a system-
atic sequence of steps designed to analyze data and derive
*Equal contribution
...
Published as a conference paper at ICLR 2025
RAPID: RETRIEVAL AUGMENTED TRAINING OF DIF-
FERENTIALLY PRIVATE DIFFUSION MODELS
Tanqiu Jiang†
Changjiang Li†
Fenglong Ma⋆
Ting Wang†
†Stony Brook University
⋆Pennsylvania State University
ABSTRACT
Differentially private diffusion models (DPDMs) harness the remarkable gener-
ative capabilities of diffusion models while enforcing differential privacy (DP)
for sensitive data.
However, existing DPDM training approaches often suffer
from significant utility loss, large memory footprint, and expensive inference
cost, impeding their practical uses. To overcome such limitations, we present
RAPID1, a novel approach that integrates retrieval augmented generation (RAG)
into DPDM training. Specifically, RAPID leverages available public data to build
a knowledge base of sample trajectories; when training the diffusion model on
private data, RAPID computes the early sampling steps as queries, retrieves sim-
ilar trajectories from the knowledge base as surrogates, and focuses on training
the later sampling steps in a differentially private manner. Extensive evaluation
using benchmark datasets and models demonstrates that, with the same privacy
guarantee, RAPID significantly outperforms state-of-the-art approaches by large
margins in generative quality, memory footprint, and inference cost, suggest-
ing that retrieval-augmented DP training represents a promising direction for de-
veloping future privacy-preserving generative models. The code is available at:
https://github.com/TanqiuJiang/RAPID.
1
INTRODUCTION
The recent advances in diffusion models have led to unprecedented capabilities of generating high-
quality, multi-modal data (Ho et al., 2020; Kong et al., 2020; Bar-Tal et al., 2024). However, training
performant diffusion models often requires massive amounts of training data, raising severe privacy
concerns in domains wherein data is sensitive. For instance, Carlini et al. (2023) show that compared
with other generative models,...
CursorCore: Assist Programming through Aligning Anything
Hao Jiang 1 Qi Liu 1 2 Rui Li 1 Shengyu Ye 1 Shijin Wang 3
Abstract
Large language models have been successfully
applied to programming assistance tasks, such
as code completion, code insertion, and instruc-
tional code editing. However, these applications
remain insufficiently automated and struggle to
effectively integrate various types of information
during the programming process, including cod-
ing history, code context, and user instructions. In
this work, we propose a new framework that com-
prehensively integrates these information sources,
and collect data to train models and evaluate their
performance. Firstly, to thoroughly evaluate how
well models align with different types of informa-
tion and the quality of their outputs, we introduce
a new benchmark, APEval (Assist Programming
Eval), to comprehensively assess the performance
of models in programming assistance tasks. Then,
for data collection, we develop a data generation
pipeline, Programming-Instruct, which synthe-
sizes training data from diverse sources, such as
GitHub and online judge platforms. This pipeline
can automatically generate various types of mes-
sages throughout the programming process. Fi-
nally, using this pipeline, we generate 219K sam-
ples, fine-tune multiple models, and develop the
CursorCore series. We show that CursorCore out-
performs other models of comparable size. This
framework unifies applications such as inline chat
and automated editing, contributes to the advance-
ment of coding assistants.
1. Introduction
Since the rise of large language models (LLMs), AI-assisted
programming technology has developed rapidly, with many
powerful LLMs being applied in this field (Zan et al., 2022;
1State Key Laboratory of Cognitive Intelligence, University of
Science and Technology of China 2Institute of Artificial Intelli-
gence, Hefei Comprehensive National Science Center 3iFLYTEK
Co., Ltd. Correspondence to: Qi Liu <qiliuql@us...
Censor Dependent Variational Inference
Chuanhui Liu 1 Xiao Wang 1
Abstract
This paper provides a comprehensive analysis
of variational inference in latent variable models
for survival analysis, emphasizing the distinctive
challenges associated with applying variational
methods to survival data. We identify a critical
weakness in the existing methodology, demon-
strating how a poorly designed variational distri-
bution may hinder the objective of survival anal-
ysis tasks—modeling time-to-event distributions.
We prove that the optimal variational distribu-
tion, which perfectly bounds the log-likelihood,
may depend on the censoring mechanism. To
address this issue, we propose censor-dependent
variational inference (CDVI), tailored for latent
variable models in survival analysis. More prac-
tically, we introduce CD-CVAE, a V-structure
Variational Autoencoder (VAE) designed for the
scalable implementation of CDVI. Further discus-
sion extends some existing theories and training
techniques to survival analysis. Extensive experi-
ments validate our analysis and demonstrate sig-
nificant improvements in the estimation of indi-
vidual survival distributions. Codes can be found
at https://github.com/ChuanhuiLiu/CDVI.
1. Introduction
Survival analysis, a fundamental topic in statistics, finds
wide-ranging applications across healthcare, insurance, qual-
ity management, and finance (Harrell et al., 2001; Nelson,
2005; Frees, 2009). It focuses on modeling the relation-
ship between time-to-event outcomes and individual demo-
graphic covariates, where the event of interest could be
death, disease progression, or similar occurrences. A key
challenge in survival analysis arises from censored obser-
vations, which provide only partial information about the
survival time, necessitating specialized methods to handle
1Department of Statistics, Purdue University, USA. Correspon-
dence to: Xiao Wang <wangxiao@purdue.edu>.
Proceedings of the 42 nd International Conference on Machine
Lea...
What Can Large Language Models Do for Sustainable Food?
Anna T. Thomas 1 Adam Yee 2 Andrew Mayne 1 Maya B. Mathur 1 Dan Jurafsky 1 Kristina Gligori´c 1
Abstract
Food systems are responsible for a third of human-
caused greenhouse gas emissions. We investigate
what Large Language Models (LLMs) can con-
tribute to reducing the environmental impacts of
food production. We define a typology of design
and prediction tasks based on the sustainable food
literature and collaboration with domain experts,
and evaluate six LLMs on four tasks in our typol-
ogy. For example, for a sustainable protein design
task, food science experts estimated that collabo-
ration with an LLM can reduce time spent by 45%
on average, compared to 22% for collaboration
with another expert human food scientist. How-
ever, for a sustainable menu design task, LLMs
produce suboptimal solutions when instructed to
consider both human satisfaction and climate im-
pacts. We propose a general framework for in-
tegrating LLMs with combinatorial optimization
to improve reasoning capabilities. Our approach
decreases emissions of food choices by 79% in
a hypothetical restaurant while maintaining par-
ticipants’ satisfaction with their set of choices.
Our results demonstrate LLMs’ potential, sup-
ported by optimization techniques, to accelerate
sustainable food development and adoption. Our
code is available at https://github.com/
thomasat/llms-sustainable-food.
1. Introduction
Global food production practices are a leading contributor to
climate change, deforestation, water pollution, biodiversity
loss, antibiotic resistance, and zoonotic disease transmis-
sion (Allan et al., 2023; Hopwood et al., 2020; Scherer et al.,
2019). Given the disproportionate harms of animal agricul-
ture in particular, which accounts for 57% of food systems
emissions, a number of companies have formed to develop
1Stanford University 2Umai Works. Correspondence to: Anna
T. Thomas <thomasat@stanford.edu>, Kristina Gligori´c <glig-
oric...
Published as a conference paper at ICLR 2025
UNI2DET: UNIFIED AND UNIVERSAL FRAMEWORK FOR
PROMPT-GUIDED MULTI-DATASET 3D DETECTION
Yubin Wang1∗, Zhikang Zou2∗, Xiaoqing Ye2, Xiao Tan2, Errui Ding2, Cairong Zhao1†
1School of Computer Science and Technology, Tongji University, 2Baidu Inc.
ABSTRACT
We present Uni2Det, a brand new framework for unified and universal multi-
dataset training on 3D detection, enabling robust performance across diverse
domains and generalization to unseen domains. Due to substantial disparities
in data distribution and variations in taxonomy across diverse domains, train-
ing such a detector by simply merging datasets poses a significant challenge.
Motivated by this observation, we introduce multi-stage prompting modules for
multi-dataset 3D detection, which leverages prompts based on the characteristics
of corresponding datasets to mitigate existing differences. This elegant design
facilitates seamless plug-and-play integration within various advanced 3D detection
frameworks in a unified manner, while also allowing straightforward adaptation
for universal applicability across datasets. Experiments are conducted across mul-
tiple dataset consolidation scenarios involving KITTI, Waymo, and nuScenes,
demonstrating that our Uni2Det outperforms existing methods by a large margin in
multi-dataset training. Notably, results on zero-shot cross-dataset transfer validate
the generalization capability of our proposed method. Our code is available at
https://github.com/ThomasWangY/Uni2Det.
1
INTRODUCTION
With the ability to capture precise geometric information of entire scenes, LiDAR has become
an essential sensor for most autonomous vehicles. Due to the rapid development of large-scale
annotated 3D LiDAR datasets such as Waymo (Sun et al., 2020), nuScenes (Caesar et al., 2020), and
KITTI (Geiger et al., 2012), LiDAR-based models play a significant role in various critical perception
tasks for autonomous vehicles, particularly in 3D object detection....
Subobject-level Image Tokenization
Delong Chen 1 2 Samuel Cahyawijaya 2 Jianfeng Liu 3 Baoyuan Wang 4 Pascale Fung 1 2
Abstract
Patch-based image tokenization ignores the mor-
phology of the visual world, limiting effective
and efficient learning of image understanding. In-
spired by subword tokenization, we introduce sub-
object-level adaptive token segmentation and ex-
plore several approaches, including superpixel,
SAM, and a proposed Efficient and PanOptiC
(EPOC) image tokenizer. Our EPOC combines
boundary detection–a simple task that can be han-
dled well by a compact model–with watershed
segmentation, which inherently guarantees no pix-
els are left unsegmented. Intrinsic evaluations
across 5 datasets demonstrate that EPOC’s seg-
mentation aligns well with human annotations
of both object- and part-level visual morphol-
ogy, producing more monosemantic tokens and
offering substantial efficiency advantages. For
extrinsic evaluation, we designed a token embed-
ding that handles arbitrary-shaped tokens, and
trained VLMs with different tokenizers on 4
datasets of object recognition and detailed cap-
tioning.
The results reveal that subobject to-
kenization enables faster convergence and bet-
ter generalization while using fewer visual to-
kens.
Project website:
https://github.
com/ChenDelong1999/subobjects.
1. Introduction
Partitioning the raw observational data stream into a man-
ageable number of segments is often beneficial for learning,
particularly when nearby elements in the data stream exhibit
strong correlations and therefore high redundancy (Cover,
1999; Barlow et al., 1961; Biederman, 1987; Marr, 2010;
Del´etang et al., 2024). In Computer Vision (CV), nearby
pixels are often highly correlated (He et al., 2022; Simon-
1Meta FAIR Paris
2The Hong Kong University of Sci-
ence and Technology 3Alibaba Group 4Zillow. Correspondence
to: Pascale Fung <pascale@ece.ust.hk>, Delong Chen <de-
long.chen@connect.ust.hk>.
Proceedings of the 42 nd International Conferenc...
Investigating and Mitigating Catastrophic Forgetting
in Medical Knowledge Injection through Internal
Knowledge Augmentation Learning
Yuxuan Zhou1, Xien Liu∗1, Xiao Zhang1, Chen Ning1, Shijin Wang2, Guoping Hu2, Ji Wu1,3,4
1Department of Electronic Engineering, Tsinghua University, Beijing, China
2iFLYTEK Research, Hefei, China
3College of AI, Tsinghua University, Beijing, China
4Beijing National Research Center for Information Science and Technology, Beijing, China
Abstract
Large Language Models (LLMs) are expected to possess comprehensive medical
knowledge to support real-world clinical applications. While domain-specific
fine-tuning effectively injects medical knowledge into LLMs, it often causes catas-
trophic forgetting of previously acquired knowledge and instruction-following ca-
pabilities. In this paper, we investigate this issue and reveal a pattern of proximity-
dependent forgetting: knowledge that is semantically or topically close to the
injected content is more likely to be forgotten, while unrelated knowledge shows
minimal degradation. Moreover, we observe that existing mitigation techniques
fail to address this type of forgetting effectively. Motivated by this observation
and inspired by human learning mechanisms, we propose InternAL (Internal
Knowledge Augmentation Learning), a novel approach that leverages LLMs’ own
internal knowledge to mitigate forgetting. InternAL first probes internal knowledge
closely related to the injection by prompting the model with questions derived
from the injected knowledge. This knowledge is then used to augment the original
injection dataset, guiding the model to retain related prior knowledge during train-
ing. Experimental results on multiple LLMs (LLaMA, Qwen) demonstrate that
InternAL significantly mitigates proximity-related forgetting while maintaining
strong knowledge injection performance. Our findings provide new insights into
the nature of catastrophic forgetting in medical knowledge injection and highlight
...
Magical: Medical Lay Language Generation via
Semantic Invariance and Layperson-tailored
Adaptation
Weibin Liao♣♠†, Tianlong Wang♣†, Yinghao Zhu♢,
Yasha Wang♣♠∗, Junyi Gao♡§∗, Liantao Ma♣♠∗
♣National Engineering Research Center For Software Engineering, Peking University
♠School of Computer Science, Peking University
♢School of Computing and Data Science, The University of Hong Kong
♡Centre for Medical Informatics, University of Edinburgh
§Health Data Research UK
{liaoweibin, tianlong.wang}@stu.pku.edu.cn, malt@pku.edu.cn
https://github.com/tianlwang/Magical.git
Abstract
Medical Lay Language Generation (MLLG) plays a vital role in improving the
accessibility of complex scientific content for broader audiences. Recent literature
to MLLG commonly employ parameter-efficient fine-tuning methods such as Low-
Rank Adaptation (LoRA) to fine-tuning large language models (LLMs) using
paired expert-lay language datasets. However, LoRA struggles with the challenges
posed by multi-source heterogeneous MLLG datasets. Specifically, through a
series of exploratory experiments, we reveal that standard LoRA fail to meet the
requirement for semantic fidelity and diverse lay-style generation in MLLG task. To
address these limitations, we propose Magical, an asymmetric LoRA architecture
tailored for MLLG under heterogeneous data scenarios. Magical employs a shared
matrix A for abstractive summarization, along with multiple isolated matrices B for
diverse lay-style generation. To preserve semantic fidelity during the lay language
generation process, Magical introduces a Semantic Invariance Constraint
to mitigate semantic subspace shifts on matrix A. Furthermore, to better adapt to
diverse lay-style generation, Magical incorporates the Recommendation-guided
Switch, an externally interface to prompt the LLM to switch between different
matrices B. Experimental results on three real-world lay language generation
datasets demonstrate that Magical consistently outperforms prompt-based methods,...
How Much Can We Forget about Data Contamination?
Sebastian Bordt 1 Suraj Srinivas 2 Valentyn Boreiko 1 Ulrike von Luxburg 1
Abstract
The leakage of benchmark data into the training
data has emerged as a significant challenge for
evaluating the capabilities of large language mod-
els (LLMs). In this work, we challenge the com-
mon assumption that small-scale contamination
renders benchmark evaluations invalid. First, we
experimentally quantify the magnitude of bench-
mark overfitting based on scaling along three di-
mensions: The number of model parameters (up
to 1.6B), the number of times an example is seen
(up to 144), and the number of training tokens (up
to 40B). If model and data follow the Chinchilla
scaling laws, minor contamination indeed leads to
overfitting. At the same time, even 144 times of
contamination can be forgotten if the training data
is scaled beyond five times Chinchilla, a regime
characteristic of many modern LLMs. Continual
pre-training of OLMo-7B corroborates these re-
sults. Next, we study the impact of the weight
decay parameter on example forgetting, showing
that empirical forgetting occurs faster than the cu-
mulative weight decay. This allows us to gauge
the degree of example forgetting in large-scale
training runs, indicating that many LLMs, includ-
ing Llama 3 405B, have forgotten the data seen at
the beginning of training.
1. Introduction
A core principle of machine learning is that a model should
not be trained on the test set used for evaluation (Donoho,
2017). For foundation models trained on Internet-scale
data, there are increasing concerns that this principle is
violated due to the leakage of benchmark evaluation data
into the training data (Xu et al., 2024; Oren et al., 2024).
Indeed, many LLM developers have found overlap between
1University of T¨ubingen, T¨ubingen AI Center, Germany 2Bosch
Research North America & Bosch Center for Artificial Intelligence
(BCAI), Sunnyvale, USA. Correspondence to: Sebastian Bordt
<sebastian.bo...
Understanding Model Reprogramming for CLIP via Decoupling Visual Prompts
Chengyi Cai * 1 Zesheng Ye * 1 Lei Feng 2 3 Jianzhong Qi 1 Feng Liu 1
Abstract
Model reprogramming adapts pretrained models
to downstream tasks by modifying only the in-
put and output spaces. Visual reprogramming
(VR) is one instance for vision tasks that adds
a trainable noise pattern (i.e., a visual prompt)
to input images to facilitate downstream classi-
fication. The existing VR approaches for CLIP
train a single visual prompt using all descriptions
of different downstream classes. However, the
limited learning capacity may result in (1) a fail-
ure to capture diverse aspects of the descriptions
(e.g., shape, color, and texture), and (2) a possi-
ble bias toward less informative attributes that do
not help distinguish between classes. In this pa-
per, we introduce a decoupling-and-reweighting
framework. Our decoupled visual prompts (DVP)
are optimized using descriptions grouped by ex-
plicit causes (DVP-cse) or unsupervised clusters
(DVP-cls). Then, we integrate the outputs of
these visual prompts with a probabilistic reweight-
ing matrix (PRM) that measures their contribu-
tions to each downstream class. Theoretically,
DVP lowers the empirical risk bound. Exper-
imentally, DVP outperforms baselines on aver-
age across 11 downstream datasets.
Notably,
the DVP-PRM integration enables insights into
how individual visual prompts influence classifi-
cation decisions, providing a probabilistic frame-
work for understanding reprogramming.
Our
code is available at https://github.com/
tmlr-group/DecoupledVP
1. Introduction
Model reprogramming (Vinod et al., 2020; Chen, 2024;
Hung et al., 2023) is shown to be an effective approach for
*Equal contribution 1School of Computing and Information Sys-
tems, The University of Melbourne 2School of Computer Science
and Engineering, Southeast University 3Idealism Technology (Bei-
jing). Correspondence to: Feng Liu <fengliu.ml@gmail.com>.
Proceedings of the 4...