← Search

Weijia Shi

41 accepted papers

2026

GRACE: Generative Representation Learning via Contrastive Policy Optimization

ICLR 2026poster

Prevailing methods for training Large Language Models (LLMs) as text encoders rely on contrastive losses that treat the model as a black-box function, discarding its generative and reasoning capabilities in favor of static embeddings. We introduce \GRACE{} (Generative Representation Learning via Con…

Cited by 0SourcecodeScholar
2026

PhenoBrain: Phenotype-Conditioned Long-Range Communication for Multi-Modal Brain Network Analysis

ICML 2026oral

Multi-modal brain network analysis aims to predict neuropsychiatric status from functional connectomes with heterogeneous phenotypes. However, most existing methods treat phenotypes as auxiliary features and perform late fusion, implicitly assuming that the connectome representation should be learne…

Cited by 0SourceScholar
2026

SplitScaling: Adaptive Scaling for Disaggregated LLM Serving Against Traffic Bursts via DRL

IJCAI 2026

The disaggregated Prefill-Decode (PD) architecture has emerged as a prominent paradigm for efficient Large Language Model inference serving. However, resource management remains a critical challenge, particularly under the dual burstiness of real-world scenarios—characterized by volatile fluctuation

Cited by 0Scholar
2026

TrustGen: A Platform of Dynamic Benchmarking on the Trustworthiness of Generative Foundation Models

ICLR 2026poster

Generative foundation models (GenFMs), such as large language models and text-to-image systems, have demonstrated remarkable capabilities in various downstream applications. As they are increasingly deployed in high-stakes applications, assessing their trustworthiness has become both a critical nece…

Cited by 0SourceScholar
2025

BRIGHT: A Realistic and Challenging Benchmark for Reasoning-Intensive Retrieval

ICLR 2025spotlight

Existing retrieval benchmarks primarily consist of information-seeking queries (e.g., aggregated questions from search engines) where keyword or semantic-based retrieval is usually sufficient. However, many complex real-world queries require in-depth reasoning to identify relevant documents that go…

Cited by 24SourcePDFScholar
2025

Fantastic Copyrighted Beasts and How (Not) to Generate Them

ICLR 2025poster

Recent studies show that image and video generation models can be prompted to reproduce copyrighted content from their training data, raising serious legal con- cerns about copyright infringement. Copyrighted characters (e.g., Mario, Batman) present a significant challenge: at least one lawsuit has…

Cited by 12SourcePDFScholar
2025

FlexOLMo: Open Language Models for Flexible Data Use

NeurIPS 2025spotlight

We introduce FlexOLMo, a new class of language models (LMs) that supports (1) distributed training without data sharing, where different model parameters are independently trained on private datasets, and (2) data-flexible inference, where these parameters along with their associated data can be eas…

Cited by 0SourceScholar
2025

Heterogeneous Swarms: Jointly Optimizing Model Roles and Weights for Multi-LLM Systems

NeurIPS 2025poster

We propose Heterogeneous Swarms, an algorithm to design multi-LLM systems by jointly optimizing model roles and weights. We represent multi-LLM systems as directed acyclic graphs (DAGs) of LLMs with topological message passing for collaborative generation. Given a pool of LLM experts and a utility f…

Cited by 0SourceScholar
2025

LMFusion: Adapting Pretrained Language Models for Multimodal Generation

NeurIPS 2025poster

We present LMFusion, a framework for empowering pretrained text-only large language models (LLMs) with multimodal generative capabilities, enabling them to understand and generate both text and images in arbitrary sequences. LMFusion leverages existing Llama-3's weights for processing texts autoregr…

Cited by 0SourceScholar
2025

MMTEB: Massive Multilingual Text Embedding Benchmark

ICLR 2025poster

Text embeddings are typically evaluated on a narrow set of tasks, limited in terms of languages, domains, and task types. To circumvent this limitation and to provide a more comprehensive evaluation, we introduce the Massive Multilingual Text Embedding Benchmark (MMTEB) -- a large-scale community-dr…

2025

MMed-RAG: Versatile Multimodal RAG System for Medical Vision Language Models

ICLR 2025poster

Artificial Intelligence (AI) has demonstrated significant potential in healthcare, particularly in disease diagnosis and treatment planning. Recent progress in Medical Large Vision-Language Models (Med-LVLMs) has opened up new possibilities for interactive diagnostic tools. However, these models oft…

2025

MUSE: Machine Unlearning Six-Way Evaluation for Language Models

ICLR 2025poster

Language models (LMs) are trained on vast amounts of text data, which may include private and copyrighted content. Data owners may request the removal of their data from a trained model due to privacy or copyright concerns. However, exactly unlearning only these datapoints (i.e., retraining with the…

Cited by 52SourcePDFScholar
2025

OLMoE: Open Mixture-of-Experts Language Models

ICLR 2025oral

We introduce OLMoE, a fully open, state-of-the-art language model leveraging sparse Mixture-of-Experts (MoE). OLMoE-1B-7B has 7 billion (B) parameters but uses only 1B per input token. We pretrain it on 5 trillion tokens and further adapt it to create OLMoE-1B-7B-Instruct. Our models outperform all…

2025

Precise Information Control in Long-Form Text Generation

NeurIPS 2025poster

A central challenge in language models (LMs) is faithfulness hallucination: the generation of information unsubstantiated by input context. To study this problem, we propose Precise Information Control (PIC), a new task formulation that requires models to generate long-form outputs grounded in a pro…

Cited by 0SourceScholar
2025

s1: Simple test-time scaling

EMNLP 2025

Test-time scaling is a promising new approach to language modeling that uses extra test-time compute to improve performance. Recently, OpenAI’s o1 model showed this capability but did not publicly share its methodology, leading to many replication efforts. We seek the simplest approach to achieve te

2024

Detecting Pretraining Data from Large Language Models

ICLR 2024poster

Although large language models (LLMs) are widely deployed, the data used to train them is rarely disclosed. Given the incredible scale of this data, up to trillions of tokens, it is all but certain that it includes potentially problematic text such as copyrighted materials, personally identifiable i…

Cited by 287SourcePDFScholar
2024

Don’t Hallucinate, Abstain: Identifying LLM Knowledge Gaps via Multi-LLM Collaboration

ACL 2024long

Despite efforts to expand the knowledge of large language models (LLMs), knowledge gaps—missing or outdated information in LLMs—might always persist given the evolving nature of knowledge. In this work, we study approaches to identify LLM knowledge gaps and abstain from answering questions when know…

2024

Evaluating Copyright Takedown Methods for Language Models

NeurIPS 2024poster

Language models (LMs) derive their capabilities from extensive training on diverse data, including copyrighted material. These models can memorize and generate content similar to their training data, potentially risking legal issues like copyright infringement. Therefore, model creators are motivat…

Cited by 21SourcePDFScholar
2024

In-Context Pretraining: Language Modeling Beyond Document Boundaries

ICLR 2024spotlight

Language models are currently trained to predict tokens given document prefixes, enabling them to zero shot long form generation and prompting-style tasks which can be reduced to document completion. We instead present IN-CONTEXT PRETRAINING, a new approach where language models are trained on a seq…

2024

Instruction-tuned Language Models are Better Knowledge Learners

ACL 2024long

In order for large language model (LLM)-based assistants to effectively adapt to evolving information needs, it must be possible to update their factual knowledge through continued training on new data. The standard recipe for doing so involves continued pre-training on new documents followed by ins…

2024

Knowledge Card: Filling LLMs' Knowledge Gaps with Plug-in Specialized Language Models

ICLR 2024oral

By design, large language models (LLMs) are static general-purpose models, expensive to retrain or update frequently. As they are increasingly adopted for knowledge-intensive tasks, it becomes evident that these design choices lead to failures to generate factual, relevant, and up-to-date knowledge.…

2024

Lemur: Harmonizing Natural Language and Code for Language Agents

ICLR 2024spotlight

We introduce Lemur and Lemur-Chat, openly accessible language models optimized for both natural language and coding capabilities to serve as the backbone of versatile language agents. The evolution from language chat models to functional language agents demands that models not only master human inte…

2024

RA-DIT: Retrieval-Augmented Dual Instruction Tuning

ICLR 2024poster

Retrieval-augmented language models (RALMs) improve performance by accessing long-tail and up-to-date knowledge from external data stores, but are challenging to build. Existing approaches require either expensive retrieval-specific modifications to LM pre-training or use post-hoc integration of the…

Cited by 142SourcePDFScholar
2024

RECOMP: Improving Retrieval-Augmented LMs with Context Compression and Selective Augmentation

ICLR 2024poster

Retrieval-augmented language models improve language models (LMs) by retrieving documents and prepending them in-context. However, these documents, often spanning hundreds of words, make inference substantially less efficient. We propose compressing the retrieved documents into textual summaries pri…

Cited by 70SourcePDFScholar
2024

REPLUG: Retrieval-Augmented Black-Box Language Models

NAACL 2024long

We introduce REPLUG, a retrieval-augmented language modeling framework that treats the language model (LM) as a black box and augments it with a tuneable retrieval model. Unlike prior retrieval-augmented LMs that train language models with special cross-attention mechanisms to encode the retrieved t…

2024

SILO Language Models: Isolating Legal Risk In a Nonparametric Datastore

ICLR 2024spotlight

The legality of training language models (LMs) on copyrighted or otherwise restricted data is under intense debate. However, as we show, model performance significantly degrades if trained only on low-risk text (e.g., out-of-copyright books or government documents), due to its limited size and domai…

2024

Scaling Retrieval-Based Language Models with a Trillion-Token Datastore

NeurIPS 2024poster

Scaling laws with respect to the amount of training data and the number of parameters allow us to predict the cost-benefit trade-offs of pretraining language models (LMs) in different configurations. In this paper, we consider another dimension of scaling: the amount of data available at inference t…

2024

Teaching LLMs to Abstain across Languages via Multilingual Feedback

EMNLP 2024main

Multilingual LLMs often have knowledge disparities across languages, with larger gaps in under-resourced languages. Teaching LLMs to abstain in the face of knowledge gaps is thus a promising strategy to mitigate hallucinations in multilingual settings. However, previous studies on LLM abstention pri…

2024

Trusting Your Evidence: Hallucinate Less with Context-aware Decoding

NAACL 2024short

Language models (LMs) often struggle to pay enough attention to the input context, and generate texts that are unfaithful or contain hallucinations. To mitigate this issue, we present context-aware decoding (CAD), which follows a contrastive output distribution that amplifies the difference between…

2024

Visual Sketchpad: Sketching as a Visual Chain of Thought for Multimodal Language Models

NeurIPS 2024poster

Humans draw to facilitate reasoning: we draw auxiliary lines when solving geometry problems; we mark and circle when reasoning on maps; we use sketches to amplify our ideas and relieve our limited-capacity working memory. However, such actions are missing in current multimodal language models (LMs).…

Cited by 41SourcePDFScholar
2023

Fine-Grained Human Feedback Gives Better Rewards for Language Model Training

NeurIPS 2023spotlight

Language models (LMs) often exhibit undesirable text generation behaviors, including generating false, toxic, or irrelevant outputs. Reinforcement learning from human feedback (RLHF)---where human preference judgments on LM outputs are transformed into a learning signal---has recently shown promise…

2023

Getting MoRE out of Mixture of Language Model Reasoning Experts

EMNLP 2023long findings

While recent large language models (LLMs) improve on various question answering (QA) datasets, it remains difficult for a single model to generalize across question types that require distinct reasoning abilities. We provide empirical evidence that state-of-the-art LLMs suffer from poor generalizabi…

Cited by 0SourceScholar
2023

Nonparametric Masked Language Modeling

ACL 2023findings

Existing language models (LMs) predict tokens with a softmax over a finite vocabulary, which can make it difficult to predict rare tokens or phrases. We introduce NPM, the first nonparametric masked language model that replaces this softmax with a nonparametric distribution over every phrase in a re…

2023

One Embedder, Any Task: Instruction-Finetuned Text Embeddings

ACL 2023findings

We introduce INSTRUCTOR, a new method for computing text embeddings given task instructions: every text input is embedded together with instructions explaining the use case (e.g., task and domain descriptions). Unlike encoders from prior work that are more specialized, INSTRUCTOR is a single embedde…

2023

PromptCap: Prompt-Guided Image Captioning for VQA with GPT-3

ICCV 2023poster

Knowledge-based visual question answering (VQA) involves questions that require world knowledge beyond the image to yield the correct answer. Large language models (LMs) like GPT-3 are particularly helpful for this task because of their strong knowledge retrieval and reasoning capabilities. To enabl…

Cited by 59PDFcodeScholar
2023

Retrieval-Augmented Multimodal Language Modeling

ICML 2023poster

Recent multimodal models such as DALL-E and CM3 have achieved remarkable progress in text-to-image and image-to-text generation. However, these models store all their knowledge (e.g., the appearance of the Eiffel Tower) in the model parameters, requiring increasingly larger models and training data…

Cited by 159SourcePDFScholar
2023

RoMQA: A Benchmark for Robust, Multi-evidence, Multi-answer Question Answering

EMNLP 2023long findings

We introduce RoMQA, the first benchmark for robust, multi-evidence, multi-answer question answering (QA). RoMQA contains clusters of questions that are derived from related constraints mined from the Wikidata knowledge graph. RoMQA evaluates robustness of QA models to varying constraints by measurin…

Cited by 0SourcecodeScholar
2023

Selective Annotation Makes Language Models Better Few-Shot Learners

ICLR 2023poster

Many recent approaches to natural language tasks are built on the remarkable abilities of large language models. Large language models can perform in-context learning, where they learn a new task from a few task demonstrations, without any parameter updates. This work examines the implications of in…

2023

Toward Human Readable Prompt Tuning: Kubrick’s The Shining is a good movie, and a good prompt too?

EMNLP 2023long findings

Large language models can perform downstream tasks in a zero-shot fashion, given natural language prompts that specify the desired behavior. Such prompts are typically hand engineered, but can also be learned with gradient-based methods from labeled data. However, it is underexplored what factors ma…

Cited by 0SourceScholar
2021

DESCGEN: A Distantly Supervised Datasetfor Generating Entity Descriptions

ACL 2021long

Short textual descriptions of entities provide summaries of their key attributes and have been shown to be useful sources of background knowledge for tasks such as entity linking and question answering. However, generating entity descriptions, especially for new and long-tail entities, can be challe…