Treasure Hunt: Real-time Targeting of the Long Tail using Training-Time Markers
Daniel Dsouza, Julia Kreutzer, Adrien Morisot, Ahmet Üstün, Sara Hooker
One of the most profound challenges of modern machine learning is performingwell on the long-tail of rare and underrepresented features. Large general-purposemodels 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 casesunderrepresented in the training corpus. Relying on prompt engineering or few-shotexamples 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 waysor rely on a fixed system prompt for maintaining performance. In this work, weask: Can we optimize our training protocols to both improve controllability andperformance on underrepresented use cases at inference time? We revisit the dividebetween training and inference techniques to improve long-tail performance whileproviding users with a set of control levers the model is trained to be responsiveto. We create a detailed taxonomy of data characteristics and task provenance toexplicitly control generation attributes and implicitly condition generations atinference time. We fine-tune a base model to infer these markers automatically,which makes them optional at inference time. This principled and flexible approachyields pronounced improvements in performance on examples from the long tailof the training distribution. Overall, we observe lifts of 5.7% across all tasks.However, treasure markers are particularly effective at finding difficult to obtaingains in the long-tail. We observe relative lifts of up to 14.1% on underrepresentedtasks like CodeRepair and absolute improvements of 35.3% on length instructionfollowing evaluations.