Permuformer improves pretraining for algebraic permutation tasks

PermuFormer: Multi-Task Pretraining for Permutation Representation in Algebraic Combinatorics

Machine Learning

Summary

The paper focuses on creating a better way to train computer models that understand permutations, which are special ways of arranging items important in math. The authors designed PermuFormer, a model trained on many related tasks to learn general features of permutations. This helps the model do well on new, unseen tasks in algebraic combinatorics after some additional training. The authors also explored how the model internally solves problems, finding some answers can be read directly while others need several steps.

What this means in practice

  • For math software developers: Enable more effective pretraining of algorithms handling permutations, improving software for algebraic combinatorics tasks.
  • For symbolic computation teams: Provide a pretrained model as a starting point to speed up development of specialized tools solving complex permutation problems.

Authors

Henry Kvinge

Abstract

Diverse pretraining has been shown to be an effective method for learning reusable, domain-aware representations that provide a starting point for fine-tuning on downstream tasks. While much of the excitement in AI for math has been concentrated in the use of frontier reasoning models to solve well-specified problems through the medium of language, narrow, specialized models remain an important component of the AI for math ecosystem. In contrast to large language models, specialized models are usually trained directly on the mathematical objects themselves (e.g., graphs, sequences of numbers) rather than the textual descriptions that characterize these objects. However, the common practice of training specialists from scratch may limit their ability to develop domain-aware representations that capture the multifaceted nature of mathematics. In this paper, we describe an approach to pretraining for permutation-focused tasks in algebraic combinatorics. We introduce PermuFormer, an autoregressive transformer trained on a 2.8 billion token multi-task, multi-encoding corpus. We show that PermuFormer is an effective starting point for fine-tuning on basic tasks unseen during pretraining and more complex research-level tasks, frequently outperforming the same architecture trained from scratch, baseline MLPs, and a fine-tuned generic language model of comparable size. We also analyze some of the internal mechanisms by which PermuFormer learns to solve training tasks. For example, we show that while some tasks can be linearly decoded directly from the internal representation of the prompt, other tasks require multiple rounds of generation before the answer can be decoded.