Treasure Hunt: Real-time Targeting of the Long Tail using Training-Time Markers

TL;DR

This paper introduces a training protocol that explicitly incorporates data and task markers to improve model performance on long-tail, underrepresented use cases, offering better control and significant gains in specialized domains.

Contribution

It presents a novel training approach using explicit markers for data and task characteristics, enhancing controllability and performance on rare and underrepresented tasks.

Findings

Average 5.7% improvement in open-ended generation quality.

Over 9.1% gains in underrepresented domains.

Up to 14.1% relative lift on specialized tasks.

Abstract

One of the most profound challenges of modern machine learning is performing well on the long-tail of rare and underrepresented features. Large general-purpose models are trained for many tasks, but work best on high-frequency use cases. After training, it is hard to adapt a model to perform well on specific use cases underrepresented in the training corpus. Relying on prompt engineering or few-shot examples to maximize the output quality on a particular test case can be frustrating, as models can be highly sensitive to small changes, react in unpredicted ways or rely on a fixed system prompt for maintaining performance. In this work, we ask: "Can we optimize our training protocols to both improve controllability and performance on underrepresented use cases at inference time?" We revisit the divide between training and inference techniques to improve long-tail performance while providing users with a set of control levers the model is trained to be responsive to. We create a detailed taxonomy of data characteristics and task provenance to explicitly control generation attributes and implicitly condition generations at inference time. We fine-tune a base model to infer these markers automatically, which makes them optional at inference time. This principled and flexible approach yields pronounced improvements in performance, especially on examples from the long tail of the training distribution. While we observe an average lift of 5.7% win rates in open-ended generation quality with our markers, we see over 9.1% gains in underrepresented domains. We also observe relative lifts of up to 14.1% on underrepresented tasks like CodeRepair and absolute improvements of 35.3% on length instruction following evaluations.