AIConfPaper
← Search

Kartik Ahuja

24 accepted papers

2026

Beyond Multi-Token Prediction: Pretraining LLMs with Future Summaries

Divyat Mahajan, Sachin Goyal, Badr Youbi Idrissi, Mohammad Pezeshki, Ioannis Mitliagkas, David Lopez-Paz +1
ICLR 2026poster

Next-token prediction (NTP) has driven the success of large language models (LLMs), but it struggles with long-horizon reasoning, planning, and creative writing, with these limitations largely attributed to teacher-forced training. Multi-token prediction (MTP) partially mitigates these issues by pre…

Cited by 0Source ↗Scholar ↗
2026

Distilled Pretraining: A modern lens of Data, In-Context Learning and Test-Time Scaling

Sachin Goyal, David Lopez-Paz, Kartik Ahuja
ICLR 2026poster

In the past year, distillation has seen a renewed prominence in large language model (LLM) pretraining, exemplified by the Llama-3.2 and Gemma model families. While distillation has historically been shown to improve statistical modeling, its effects on new paradigms key to modern LLMs—such as test-…

Cited by 0Source ↗Scholar ↗
2026

Teaching Models to Teach Themselves: Reasoning at the Edge of Learnability

Shobhita Sundaram, John Quan, Ariel Kwiatkowski, Kartik Ahuja, Yann Ollivier, Julia Kempe
ICML 2026spotlight

RL methods for finetuning large reasoning models stall on datasets with low initial success rates, and thus little training signal. We investigate a fundamental question: Can a pretrained LLM leverage latent knowledge to generate an automated curriculum for problems it cannot solve? We explore this …

Cited by 0Source ↗Scholar ↗
2025

Compositional Risk Minimization

Divyat Mahajan, Mohammad Pezeshki, Charles Arnal, Ioannis Mitliagkas, Kartik Ahuja, Pascal Vincent
ICML 2025poster

Compositional generalization is a crucial step towards developing data-efficient intelligent machines that generalize in human-like ways. In this work, we tackle a challenging form of distribution shift, termed compositional shift, where some attribute combinations are completely absent at training…

Cited by 3Source ↗PDF ↗Scholar ↗
2024

Multi-Domain Causal Representation Learning via Weak Distributional Invariances

Kartik Ahuja, Amin Mansouri, Yixin Wang
AISTATS 2024poster

Causal representation learning has emerged as the center of action in causal machine learning research. In particular, multi-domain datasets present a natural opportunity for showcasing the advantages of causal representation learning over standard unsupervised representation learning. While recent…

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

Locally Invariant Explanations: Towards Stable and Unidirectional Explanations through Local Invariant Learning

Amit Dhurandhar, Karthikeyan Natesan Ramamurthy, Kartik Ahuja, Vijay Arya
NeurIPS 2023poster

Locally interpretable model agnostic explanations (LIME) method is one of the most popular methods used to explain black-box models at a per example level. Although many variants have been proposed, few provide a simple way to produce high fidelity explanations that are also stable and intuitive. In…

Cited by 8Source ↗PDF ↗Scholar ↗
2023

Model Ratatouille: Recycling Diverse Models for Out-of-Distribution Generalization

Alexandre Rame, Kartik Ahuja, Jianyu Zhang, Matthieu Cord, Leon Bottou, David Lopez-Paz
ICML 2023poster

Foundation models are redefining how AI systems are built. Practitioners now follow a standard procedure to build their machine learning solutions: from a pre-trained foundation model, they fine-tune the weights on the target task of interest. So, the Internet is swarmed by a handful of foundation m…

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

Reusable Slotwise Mechanisms

Trang Nguyen, Amin Mansouri, Kanika Madan, Nguyen Duy Khuong, Kartik Ahuja, Dianbo Liu +1
NeurIPS 2023poster

Agents with the ability to comprehend and reason about the dynamics of objects would be expected to exhibit improved robustness and generalization in novel scenarios. However, achieving this capability necessitates not only an effective scene representation but also an understanding of the mechanism…

Cited by 4Source ↗PDF ↗Scholar ↗
2023

Why does Throwing Away Data Improve Worst-Group Error?

Kamalika Chaudhuri, Kartik Ahuja, Martin Arjovsky, David Lopez-Paz
ICML 2023oral

When facing data with imbalanced classes or groups, practitioners follow an intriguing strategy to achieve best results. They throw away examples until the classes or groups are balanced in size, and then perform empirical risk minimization on the reduced training set. This opposes common wisdom in…

Cited by 18Source ↗PDF ↗Scholar ↗
2022

Finding Valid Adjustments under Non-ignorability with Minimal DAG Knowledge

Abhin Shah, Karthikeyan Shanmugam, Kartik Ahuja
AISTATS 2022poster

Treatment effect estimation from observational data is a fundamental problem in causal inference. There are two very different schools of thought that have tackled this problem. On the one hand, the Pearlian framework commonly assumes structural knowledge (provided by an expert) in the form of direc…

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

Properties from mechanisms: an equivariance perspective on identifiable representation learning

Kartik Ahuja, Jason Hartford, Yoshua Bengio
ICLR 2022spotlight

A key goal of unsupervised representation learning is ``inverting'' a data generating process to recover its latent properties. Existing work that provably achieves this goal relies on strong assumptions on relationships between the latent variables (e.g., independence conditional on auxiliary info…

Cited by 41Source ↗PDF ↗Scholar ↗
2022

Weakly Supervised Representation Learning with Sparse Perturbations

Kartik Ahuja, Jason Hartford, Yoshua Bengio
NeurIPS 2022accept

The theory of representation learning aims to build methods that provably invert the data generating process with minimal domain knowledge or any source of supervision. Most prior approaches require strong distributional assumptions on the latent variables and weak supervision (auxiliary information…

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

Adversarial Feature Desensitization

Pouya Bashivan, Reza Bayat, Adam Ibrahim, Kartik Ahuja, Mojtaba Faramarzi, Touraj Laleh +2
NeurIPS 2021poster

Neural networks are known to be vulnerable to adversarial attacks -- slight but carefully constructed perturbations of the inputs which can drastically impair the network's performance. Many defense methods have been proposed for improving robustness of deep networks by training them on adversarial…

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

Can Subnetwork Structure Be the Key to Out-of-Distribution Generalization?

Dinghuai Zhang, Kartik Ahuja, Yilun Xu, Yisen Wang, Aaron Courville
ICML 2021oral

Can models with particular structure avoid being biased towards spurious correlation in out-of-distribution (OOD) generalization? Peters et al. (2016) provides a positive answer for linear cases. In this paper, we use a functional modular probing method to analyze deep model structures under OOD set…

Cited by 109Source ↗PDF ↗Scholar ↗
2021

Conditionally independent data generation

Kartik Ahuja, Prasanna Sattigeri, Karthikeyan Shanmugam, Dennis Wei, Karthikeyan Natesan Ramamurthy, Murat Kocaoglu
UAI 2021poster

Conditional independence (CI) is a fundamental concept with wide applications in machine learning and causal inference. Although the problems of testing CI and estimating divergences have been extensively studied, the complementary problem of generating data that satisfies CI has received much less…

Cited by 5Source ↗PDF ↗Scholar ↗
2021

Empirical or Invariant Risk Minimization? A Sample Complexity Perspective

Kartik Ahuja, Jun Wang, Amit Dhurandhar, Karthikeyan Shanmugam, Kush R. Varshney
ICLR 2021poster

Recently, invariant risk minimization (IRM) was proposed as a promising solution to address out-of-distribution (OOD) generalization. However, it is unclear when IRM should be preferred over the widely-employed empirical risk minimization (ERM) framework. In this work, we analyze both these framewor…

Cited by 104Source ↗PDF ↗Scholar ↗
2021

Invariance Principle Meets Information Bottleneck for Out-of-Distribution Generalization

Kartik Ahuja, Ethan Caballero, Dinghuai Zhang, Jean-Christophe Gagnon-Audet, Yoshua Bengio, Ioannis Mitliagkas +1
NeurIPS 2021spotlight

The invariance principle from causality is at the heart of notable approaches such as invariant risk minimization (IRM) that seek to address out-of-distribution (OOD) generalization failures. Despite the promising theory, invariance principle-based approaches fail in common classification tasks, whe…

Cited by 313Source ↗PDF ↗Scholar ↗
2021

Linear Regression Games: Convergence Guarantees to Approximate Out-of-Distribution Solutions

Kartik Ahuja, Karthikeyan Shanmugam, Amit Dhurandhar
AISTATS 2021poster

Recently, invariant risk minimization (IRM) (Arjovsky et al. 2019) was proposed as a promising solution to address out-of-distribution (OOD) generalization. In Ahuja et al. (2020), it was shown that solving for the Nash equilibria of a new class of “ensemble-games” is equivalent to solving IRM. In t…

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

Treatment Effect Estimation Using Invariant Risk Minimization

Abhin Shah, Kartik Ahuja, Karthikeyan Shanmugam, Dennis Wei, Kush R. Varshney, Amit Dhurandhar
ICASSP 2021accepted

Inferring causal individual treatment effect (ITE) from observational data is a challenging problem whose difficulty is exacerbated by the presence of treatment assignment bias. In this work, we propose a new way to estimate the ITE using the domain generalization framework of invariant risk minimiz…

Cited by 0Source ↗Scholar ↗