🤖 machine learning

Learning the symmetric group: large from small

This article proposes a scalable machine learning method in which Transformer models trained to predict permutations in smaller symmetric groups (such as $S_{10}$ ) using specific transposition strategies generalize with near-perfect accuracy to significantly larger groups (such as $S_{25}$ ), employing techniques such as identity augmentation and partitioned windows to address challenges in data generation and interpretability.

Original authors: Max Petschack, Alexandr Garbali, Jan de Gier

Published 2026-05-08

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Max Petschack, Alexandr Garbali, Jan de Gier

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Idea: Teaching a Student to Solve Giant Puzzles with Tiny Practice Sets

Imagine you want to teach a student to solve a huge, complex puzzle with 25 pieces (or even 100). Normally, you would give them practice puzzles of exactly that size. But what if you only gave them practice puzzles with 10 pieces?

This paper asks: Can a computer (specifically an AI called a "Transformer") learn the rules of a giant puzzle by practicing only on small versions of it, and then successfully solve the giant version without ever having seen one before?

The answer, according to this study, is Yes. The AI learned the logic of a massive mathematical system by training on a tiny part of it, and then transferred that knowledge to handle much larger, more complex systems with near-perfect accuracy.

The Characters in Our Story

The Symmetric Group ( $S_n$ ): Imagine this as a giant game of "shuffling a deck of cards."
- If you have a deck with $n$ cards (numbered 1 to $n$ ), a "permutation" is simply a specific order of these cards.
- The "symmetric group" is the collection of every possible way you can shuffle this deck.
- The goal is to look at a list of instructions (a "word") that tells you how to shuffle the cards and predict the final order of the deck.
The Instructions (Transpositions):
- General Transpositions: Imagine you can pick any two cards in the deck and swap them. This is like a "magic wand" that can instantly swap any two objects.
- Adjacent Transpositions: Imagine you can only swap cards that are directly next to each other. This is much harder. To swap card #1 and card #10, you must shuffle them past each other one by one. This creates a much longer, more complicated list of instructions.
The AI (The Transformer): This is a type of machine learning model known for reading text and understanding patterns. Here, instead of sentences, it reads lists of mathematical instructions.

The Experiment: Two Different Challenges

The researchers conducted two main experiments to see if the AI could scale from "small" to "large."

Challenge 1: The "Magic Wand" (General Transpositions)

The Training: The AI was trained only on shuffling decks with 10 cards. It learned to follow instructions to swap any two cards in a 10-card deck.
The Test: Then, they asked the AI to solve shuffling problems for a deck with 25 cards.
The Result: The AI was correct almost 100% of the time. It didn't just memorize the rules for 10 cards; it recognized the underlying logic of "swapping" and applied it to a much larger deck it had never seen before.

Challenge 2: The "Neighbor Swap" (Adjacent Transpositions)

The Training: This was harder. The AI was trained on a 10-card deck where it was allowed to swap only neighbors.
The Problem: If you can only swap neighbors, the instructions become very long. A simple swap of the first and last card requires many steps.
The Trick (Partitioned Windows): The researchers found that the AI was getting lazy. It simply memorized the specific "window" of cards it was looking at. To fix this, they used a method with "partitioned windows." Imagine you break the long list of instructions into sections and shuffle these sections around so the AI cannot rely on position. It had to learn the actual logic of the swaps.
The Test: They tested the AI on a deck with 16 cards.
The Result: Again, the AI achieved near-100% accuracy.

How Did They Make It Work? (The Secret Recipe)

The researchers used a clever trick called "Identity Augmentation."

Imagine you are writing a recipe, but the recipe book requires every recipe to be exactly 50 steps long. Some recipes are naturally short (only 5 steps). To fit the book, you must add "dummy steps" that say "do nothing" (like "stand still for 5 seconds") until you reach 50 steps.

The AI had to learn that these "do-nothing" steps did not change the result. By filling out the short instructions with these "identity" steps, the AI learned to ignore the noise and focus on the actual mathematics.

What Did the AI Actually Learn?

The researchers looked into the AI's "brain" (its internal data representations) to see what it was doing.

It learned relationships: The AI recognized that swapping card A with card B is the same as swapping B with A.
It learned structure: It recognized that the order of swaps matters, but also learned the rules for when the order is not important.
It didn't cheat: The AI didn't just memorize the answers. It had to learn the "algorithm" of shuffling because the test questions were different from the training questions.

The Conclusion

This paper proves that AI models can learn complex mathematical rules from small examples and apply them to much larger, more complex versions of the same problem.

The Analogy: It is like teaching a child how to tie their shoes using a small practice board, and then giving them a pair of giant boots. The child, having learned the logic of the knot, can tie the giant boots perfectly.
The Limit: The authors note that while this works excellently for the "symmetric group" (shuffling cards), other mathematical groups might be messier and harder to learn. Nevertheless, this success suggests that AI could eventually help solve other difficult mathematical problems, such as the "Unknot Problem" (untangling knots), which is currently very difficult for computers.

In short: The AI learned to become a master shuffler by practicing on a small deck, proving that machines with the right training can generalize from "small" to "large" in pure mathematics.

Technical Conclusion: Learning the Symmetric Group: Large from Small

Problem Statement
The article investigates whether Transformer-based neural networks can learn the structure of the symmetric group $S_n$ by training exclusively on smaller subgroups $S_m$ (where $m < n$ ) and subsequently generalizing to predict permutations in larger groups, without hard-coding group relations. The specific task involves taking a word $w \in S_n$ , expressed as a product of transpositions, and predicting the resulting permutation of the set $\{1, \dots, n\}$ . The authors address two different formulations of this problem:

General Transpositions: Words factored using arbitrary transpositions $s_{i,j}$ .
Adjacent Transpositions: Words factored exclusively using elementary adjacent transpositions $s_i = s_{i,i+1}$ .

A central challenge lies in handling variable word lengths. The authors note that the lengths of reduced words vary from $0$ (identity) to $n-1$ (general case) or $n(n-1)/2$ (adjacent case). To address this, models are trained on fixed-length, non-reduced words padded with identity transpositions.

Methodology
The authors employ an autoregressive Transformer architecture (Figure 1) that predicts permutation tokens sequentially. The input consists of a "context window" containing the transposition sequence and the partially predicted permutation.

Data Generation and Augmentation:
- Identity Padding: To handle variable lengths, words are padded to a maximum length $N$ with transpositions that act as identities under group relations. The model must implicitly learn these relations to ignore the padding.
- Training on Subgroups: The model is trained on data generated from $S_m$ (specifically $m=10$ ) but tested on $S_n$ ( $n=25$ for the general case, $n=16$ for the adjacent case).
- Renaming Strategy: To generate training data for $S_n$ $S_{n}$ that includes only $S_m$ $S_{m}$ , the authors map words from $S_m$ $S_{m}$ to $S_n$ $S_{n}$ via a random permutation $\sigma \in S_n$ $σ \in S_{n}$ .
  - For general transpositions, a simple random renaming of indices $\{1, \dots, m\}$ to a random subset of $\{1, \dots, n\}$ is used.
  - For adjacent transpositions, a naive "window" method (renaming $s_1, \dots, s_{m-1}$ to a contiguous block $s_{k+1}, \dots$ ) failed, as the model learned to identify the window position rather than the group structure. To counteract this, the authors introduced a partitioned window method. This involves decomposing $m$ into a composition (e.g., $12 = 3+6+3$ ), creating multiple non-contiguous windows of varying lengths within $S_n$ , and mapping the input tokens to these windows with random shifts. This forces the model to learn local group relations rather than global position indicators.
Architecture Details:
- Tokenization: General transpositions are mapped to integers $x_k = i_k - 1 + n(j_k - 1)$ . Adjacent transpositions are mapped directly to indices $x_k = i_k$ .
- Embeddings: The model uses learned token embeddings and positional embeddings. The input is processed through standard Transformer layers (Multi-Head Attention, Feed-Forward, Normalization) before a linear layer and softmax output probabilities follow.

Main Results
Experiments show that the models achieve nearly 100% accuracy on Out-of-Distribution (OOD) test sets:

General Transpositions: A model trained on $S_{10}$ generalized to $S_{25}$ with zero training loss and near-perfect test performance (Figure 2).
Adjacent Transpositions: A model trained on $S_{10}$ generalized to $S_{16}$ with similarly high performance (Figure 3). The authors note that the adjacent case is theoretically more complex due to longer word lengths and stricter locality constraints; nevertheless, the partitioned window method successfully enabled generalization.

Observations on Interpretability
The authors analyzed the self-similarity heatmaps of the learned embeddings:

General Transpositions: The token embedding matrix revealed a regular grid structure, suggesting the model learned high-level relationships between generators (e.g., $s_{i,j}$ is strongly related to $s_{i,k}$ and $s_{j,i}$ ). The positional embedding showed no intrinsic ordering under general transpositions.
Adjacent Transpositions: The positional embedding for adjacent transpositions showed a surprising block substructure within the first $n(n-1)/2$ rows, although the authors note this was not consistent across all group sizes and its significance remains unclear.

Significance and Claims
The article claims to demonstrate that Transformers can learn the underlying algorithmic structure of group multiplication from small training datasets and generalize to significantly larger groups, even when the input representation changes (via renaming and partitioning).

Scalability: The work proposes a scalable method where models trained on simpler structures ( $S_m$ ) can solve tasks for larger, more complex groups ( $S_n$ ).
Out-of-Distribution Learning: This serves as a specific example of OOD learning in mathematics, showing that features learned on small subgroups can be universal enough to apply to larger groups.
Future Directions: The authors suggest this approach could provide insights into the word problem for finitely generated groups and applications to NP-hard problems in comparative genomics. They specifically identify the braid group as a target for future work, motivated by its relevance to the unknot problem, and note that the symmetric group is "tame," while more exotic groups will present greater challenges.

The authors remain modest regarding the internal mechanisms, noting that it is an "open question" whether the models discovered a general algorithm or learned specific representations, and they plan to conduct a mechanistic interpretability analysis in future work.