Computational Algebra with Attention: Transformer Oracles for Border Basis Algorithms

📄 Paper  |  📝 Blog

Authors: Hiroshi Kera*, Nico Pelleriti*, Yuki Ishihara, Max Zimmer, Sebastian Pokutta (*equal contribution)

This repository accompanies our NeurIPS 2025 paper "Computational Algebra with Attention: Transformer Oracles for Border Basis Algorithms".

Polynomial system solving appears in cryptography, robotics, and optimization, but classical methods often spend a lot of time on intermediate computations that do not end up in the final basis.

OBBA (Oracle Border Basis Algorithm) uses a Transformer to prioritize useful expansions and keeps a verified fallback path to preserve correctness. In our benchmarks, this gives up to 3.5x speedups.

Highlights

This section summarizes the main idea, runtime impact, and model design choices in three quick figures.

1) Why oracle guidance helps

Border basis construction expands candidates iteratively. The key difference is whether expansion is in all directions (classical) or targeted (oracle-guided).

Figure 1:(a) shows a border basis toy example with order ideal ${1, y}$ and border terms ${y^2, xy, x}$. (b) shows naive BBA expansion, where many candidates are explored although several are ultimately unnecessary. (c) shows oracle-guided expansion, where only likely-useful candidates are expanded, reaching the same target with fewer steps.

2) What this buys in runtime

The plot below compares runtime scaling across finite fields and degrees for classical and oracle-guided variants.

Runtime barplots for $\mathbb{F}7$, $\mathbb{F}{127}$, and $\mathbb{F}_{31}$ across total degree $d \in {2,3,4}$. Bars compare Border Basis Algorithm (BBA) and oracle/improved variants (OBBA, IBBA+FGE, OBBA+FGE), with IBBA included as an additional baseline.

3) How we made the Transformer practical

A key part of OBBA is how polynomial states are represented for the model. Instead of long infix token streams, we use a monomial-level embedding scheme that keeps algebraic structure but uses far fewer tokens.

Token count comparison across variable counts. The monomial representation is much more compact than infix baselines, reducing sequence length by orders of magnitude in larger settings. This yields better scaling for self-attention (about $\mathcal{O}(n)$ fewer tokens and $\mathcal{O}(n^2)$ less attention memory for $n$ variables) while improving prediction quality.

Environment setup

SageMath 10.0+ cannot currently be installed via apt-get + pip install alone (noted May 27, 2025). Use the instructions on this page.

Reproduce the paper results

Main experiments

Generate datasets:

bash sh/generate_dataset_sweep.sh

Train models:

bash sh/train_sweep.sh

Reproduce Table 1:

python scripts/analysis/model_evaluation.py

For Table 2, first train/load the model for field $\mathbb{F}_{31}$ with 5 leading terms and set save_path in sweeps/n=5/evaluation.yaml. Then run:

wandb sweep sweeps/n=5/evaluation.yaml

For Figure 3, set save_path and run the sweep defined in sweeps/n=4/ood_evaluation.yaml.

Additional experiment

Infix vs monomial representation (Table 3):

bash train_infix_vs_monomial.sh

Token-count analysis used for Figure 7:

python scripts/analysis/infix_vs_monomial_token_count_and_plot.py

General workflow

Dataset generation:

1. Prepare a problem config in config/problems/.
2. Prepare and run sh/generate_getaset.sh.

bash sh/generate_getaset.sh

Training:

1. Prepare an experiment config in config/experiments/.
2. Prepare and run sh/train.sh.

bash sh/train.sh

--dryrun runs a smaller experiment and logs to the dryrun W&B project.

OBBA evaluation:

1. Set save_path in the relevant sweep file (for example, sweeps/n=5/evaluation.yaml).
2. Launch the sweep and agent with W&B.

Project structure

Transformer-BB/
├── config/         # Configuration files
├── data/           # Dataset and preprocessing scripts
├── figs/           # Figures and plots
├── notebook/       # Jupyter notebooks
├── results/        # Experimental results
├── scripts/        # Script files
├── sh/             # Shell scripts
├── src/            # Source code
├── sweeps/         # W&B yaml sweep files
└── tests/          # Test cases and unit tests

Citation

If you use this code in your research, please cite:

@misc{kera_pelleriti2025computational,
  title={Computational Algebra with Attention: Transformer Oracles for Border Basis Algorithms},
  author={Hiroshi Kera and Nico Pelleriti and Yuki Ishihara and Max Zimmer and Sebastian Pokutta},
  year={2025},
  archivePrefix={arXiv},
  eprint={2505.23696}
}