Title: Teaching Algorithms to Small Language Models
Abstract: Recent studies have shown that large language models can perform surprising algorithmic reasoning tasks by leveraging chain-of-thought prompting and test-time computation. Yet they often struggle to generalize beyond the input length and difficulty of simpler tasks that they saw during training. In this talk, we examine the mechanisms of these capabilities by training small transformers—initialized at random—to perform arithmetic operations (e.g., addition and multiplication), string manipulation, and maze solving using the next-token prediction objective. We find that simple formatting changes, chain-of-thought data, and iterative test-time computation (“looping”) significantly improve test accuracy, sample complexity, and convergence speed. Despite these gains, these models still exhibit persistent failures in length and hardn ess generalization; they do not seem to learn the underlying algorithms, and learn shortcuts that only fit the format of the training distribution. While previous work on length generalization has often relied on departing from the base transformer architecture, we propose a simpler and more general solution that works for standard transformers, like GPT and Llama models. Specifically, we use a self-improvement process where models begin by solving easier problems, then use their own outputs to generate training data for progressively harder tasks. Scaling this weak-to-strong training approach yields (seemingly) unbounded improvements in both length and hardness generalization, allowing models to solve problem instances far exceeding the difficulty of those in the training data distribution. We notice that “controlled sampling” of problem difficulty is key; without it, generalization performance plateaus. Our results indicate that with careful self-supervision small transformers can transcend superficial pattern matching failures and develop genuine multi-step reasoning skills.
Bio: Dimitris Papailiopoulos is a Principal Researcher at Microsoft Research leading a language model reasoning effort, and is also the Jay & Cynthia Ihlenfeld Associate Professor of ECE at the University of Wisconsin-Madison. His research interests span machine learning, large-scale systems, and information theory, with a current focus on understanding the intricacies and improving the performance of large-language models. Before joining MSR and UW-Madison, Dimitris was a postdoctoral researcher at UC Berkeley and a member of the AMPLab. He spent three wonderful years at USC as a PhD, and then earned his degree from UT Austin, under the guidance of Alex Dimakis. He received his ECE Diploma M.Sc. degree from the Technical University of Crete, in Greece. Dimitris is a recipient of the NSF CAREER Award (2019), three years of Sony Faculty Innovation Awards (2018, 2019 and 2020), a joint IEEE ComSoc/ITSoc Best Paper Award (2020), an IEEE Signal Processing Society, Young Author Best Paper Award (2015), the Vilas Associate Award (2021), the Emil Steiger Distinguished Teaching Award (2021), and the Benjamin Smith Reynolds Award for Excellence in Teaching (2019). In 2018, he co-founded MLSys, a new conference that targets research at the intersection of machine learning and systems.