Author name cluster

Benjamin Cohen-Wang

Possible papers associated with this exact author name in Arrow. This page groups case-insensitive exact name matches and is not a full identity disambiguation profile.

3 papers

2 author rows

TMLR Journal 2025 Journal Article

Ask Your Distribution Shift if Pre-Training is Right for You

Benjamin Cohen-Wang
Joshua Vendrow
Aleksander Madry

Pre-training is a widely used approach to develop models that are robust to distribution shifts. However, in practice, its effectiveness varies: fine-tuning a pre-trained model improves robustness significantly in some cases but *not at all* in others (compared to training from scratch). In this work, we seek to characterize the failure modes that pre-training *can* and *cannot* address. In particular, we focus on two possible failure modes of models under distribution shift: poor extrapolation (e.g., they cannot generalize to a different domain) and biases in the training data (e.g., they rely on spurious features). Our study suggests that, as a rule of thumb, pre-training can help mitigate poor extrapolation but not dataset biases. After providing theoretical motivation and empirical evidence for this finding, we explore two of its implications for developing robust models: (1) pre-training and interventions designed to prevent exploiting biases have complementary robustness benefits, and (2) fine-tuning on a (very) small, non-diverse but *de-biased* dataset can result in significantly more robust models than fine-tuning on a large and diverse but biased dataset.

PDF Details

ICML Conference 2025 Conference Paper

SelfCite: Self-Supervised Alignment for Context Attribution in Large Language Models

Yung-Sung Chuang
Benjamin Cohen-Wang
Zejiang Shen 0001
Zhaofeng Wu
Hu Xu 0001
Xi Victoria Lin
James R. Glass
Shang-Wen Li 0001

We introduce SelfCite, a novel self-supervised approach that aligns LLMs to generate high-quality, fine-grained, sentence-level citations for the statements in their generated responses. Instead of only relying on costly and labor-intensive annotations, SelfCite leverages a reward signal provided by the LLM itself through context ablation: If a citation is necessary, removing the cited text from the context should prevent the same response; if sufficient, retaining the cited text alone should preserve the same response. This reward can guide the inference-time best-of-N sampling strategy to improve citation quality significantly, as well as be used in preference optimization to directly fine-tune the models for generating better citations. The effectiveness of SelfCite is demonstrated by increasing citation F1 up to 5. 3 points on the LongBench-Cite benchmark across five long-form question answering tasks. The source code is available at https: //github. com/facebookresearch/SelfCite.

Details

NeurIPS Conference 2024 Conference Paper

ContextCite: Attributing Model Generation to Context

Benjamin Cohen-Wang
Harshay Shah
Kristian Georgiev
Aleksander Mądry

How do language models use information provided as context when generating a response? Can we infer whether a particular generated statement is actually grounded in the context, a misinterpretation, or fabricated? To help answer these questions, we introduce the problem of context attribution: pinpointing the parts of the context (if any) that led a model to generate a particular statement. We then present ContextCite, a simple and scalable method for context attribution that can be applied on top of any existing language model. Finally, we showcase the utility of ContextCite through three applications: (1) helping verify generated statements(2) improving response quality by pruning the context and(3) detecting poisoning attacks. We provide code for ContextCite at https: //github. com/MadryLab/context-cite.

PDF Details DOI