Search papers, labs, and topics across Lattice.
This study explores how small transformers learn modular integer multiplication, a non-invertible operation, by introducing the monoid extension, which generalizes Group Composition via Representation (GCR). The researchers demonstrate that transformers can effectively partition input space into local hierarchical algebraic regions, allowing for the application of Fourier mechanisms, which leads to organized embeddings and class-sensitive attention routing. Key findings reveal that local character features significantly influence the model's output logits, indicating that representation-theoretic mechanisms can extend beyond traditional group operations.
Transformers can learn modular multiplication by partitioning input space into local algebraic regions, revealing surprising new insights into their reasoning capabilities.
Transformers have demonstrated a remarkable ability to learn algorithmic reasoning, yet mechanistic analyses have mostly focused on globally invertible operations such as cyclic addition and group composition. In this work, we investigate how small transformers learn modular integer multiplication over composite moduli, a fundamentally non-invertible operation due to the presence of zero-divisors. We propose the monoid extension: a localized generalization of Group Composition via Representation (GCR) that suggests the learned computation does not rely on a single global representation space. Instead, the model partitions the input space into local hierarchical algebraic regions, where group-like structure survives and Fourier mechanisms can be applied. In transformers trained on square-free modular multiplication, we find that embeddings organize around these regions, attention exhibits class-sensitive routing and low-rank write directions, and local character features explain a large fraction of the model's output logits. Our results suggest that representation-theoretic mechanisms previously identified for group operations can extend beyond groups to more general structures.
