hep-thMar 11, 2026

Learning to Unscramble: Simplifying Symbolic Expressions via Self-Supervised Oracle Trajectories

Authors: David Shih

Organizations: NHETC, Dept. of Physics and Astronomy, Rutgers University, Piscataway, NJ 08854, USA

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

We present a new self-supervised machine learning approach for symbolic simplification of complex mathematical expressions. Training data is generated by scrambling simple expressions and recording the inverse operations, creating oracle trajectories that provide both goal states and explicit paths to reach them. A permutation-equivariant, transformer-based policy network is then trained on this data step-wise to predict the oracle action given the input expression. We demonstrate this approach on two problems in high-energy physics: dilogarithm reduction and spinor-helicity scattering amplitude simplification. In both cases, our trained policy network achieves near perfect solve rates across a wide range of difficulty levels, substantially outperforming prior approaches based on reinforcement learning and end-to-end regression. When combined with contrastive grouping and beam search, our model achieves a 100% full simplification rate on a representative selection of 5-point gluon tree-level amplitudes in Yang-Mills theory, including expressions with over 200 initial terms.

Figures & tables

Figure 1: Architecture of the policy network for symbolic simplification. Each term’s feature vector is embedded and processed by a Transformer encoder with a prepended learnable [CLS] token. No positional encoding is used, respecting the permutation symmetry of terms. The permutation equivariant policy head takes the output of the transformer and returns probabilities (softmax) over the action space.
  Our modelDSZ
 NNSource-relTarget-rel 
By target terms
0810011
19021241
21,16012104
31,86501253
Table 1: Dilogarithm simplification failures on the clean test set of DSZ (4,737 samples after excluding 164 overlapping with our training data and 128 pathological, with 29 in both), broken down by number of initial terms and by the number of scrambles. Source-relative: failure to produce a correct output with strictly fewer terms than the initial scrambled expression. Target-relative: failure to reach term count of the initial unscrambled expression. For the seq2seq model of DSZ, we only report one column of failure numbers as there is no empirical difference between the two criteria: it turns out that their model either outputs expressions that exactly reproduce the target, or it outputs incorrect outputs that are not mathematically equivalent to the input.
Figure 2: Solve rate vs. scramble depth for dilogarithm simplification. Our model (blue) maintains near-100% performance under both source-relative and target-relative criteria, even beyond the training range (shaded). The seq2seq model of DSZ (orange) degrades at higher scramble depths.
Figure 3: Average number of steps to solve vs. scramble depth for dilogarithm simplification, with ±\pm 1 standard deviation error bars. The dashed line marks y=xy=x (steps equal to scramble depth); the shaded region indicates the training range (scramble depths 1–7). The model consistently finds shorter paths than the scramble depth, indicating it learns to bypass redundancy introduced by the scrambling process.
  Our modelCDS
 NNSource-relTarget-relSource-relTarget-rel
4-point, by target terms
01,3380103
14,77500932
22,366122853
31,390288288
Table 2: Scattering amplitude simplification failures on the test sets of CDS, excluding pathological samples where the target is not strictly simpler than the source (123 for 4pt, 5 for 5pt, 0 for 6pt) and train/test overlap samples (28 for 4pt, 9 for 5pt, 0 for 6pt), broken down by number of target terms. Bracket counts are power-aware ( ⟨ij⟩n\langle ij\rangle^{n} counts as nn ). Both criteria are applied uniformly to both models. Source-relative: failure to produce output strictly simpler in at least one of brackets or terms (without increasing the other). Target-relative: failure to produce output with both bracket count and term count ≤\leq the known target. For CDS, success requires any valid beam hypothesis ( B=20B{=}20 ) to satisfy the criterion.
Figure 4: Solve rate vs. number of target terms for 4-point (left), 5-point (center), and 6-point (right) amplitudes. Performance is shown under source-relative (dashed) and target-relative (solid) criteria for our model (blue) and CDS with B=20B{=}20 (red). Gray bars show the number of test samples at each target term count.
Figure 5: Solve rate vs. source bracket count for 4-point (left), 5-point (center), and 6-point (right) amplitudes. Our model (blue) maintains near-100% solve rates across all starting complexities under both source-relative (dashed) and target-relative (solid) criteria. The model of CDS with beam size B=20B{=}20 (red) degrades with increasing bracket count. Gray bars show the number of test samples at each bracket count.
Figure 6: Solve rate for 5-point Yang-Mills partial amplitudes as a function of the initial number of numerator terms, binned in groups of 25. Stacked bars show our results: blue indicates forms solved by the greedy contrastive-grouping phase alone, green indicates forms additionally requiring beam search. The orange curve shows the CDS sequential simplification solve rate, digitized from Fig. 9 (right panel) of Ref. [ 9 ] and averaged within each bin. Numbers above each bar indicate the number of forms in that bin. Our method achieves a 100% solve rate across all bins (103/103 forms), while CDS performance degrades with expression complexity.

Explore similar work

CardsList

  1. Machine learning modularity

    Jan 5, 2026Yi Fan, Vishnu Jejjala, Yang LeiSymbolic Reasoning

    School of Physical Science and Technology, Soochow University, 333 Ganjiang Road, Suzhou 215006, P.R. China · Institute for Advanced Study, Soochow University, 333 Ganjiang Road, Suzhou 215006, P.R. China · Mandelstam Institute for Theoretical Physics, School of Physics and NITheCS, University of the Witwatersrand, Johannesburg 2050, South Africa +1

  2. Learning to Unscramble Feynman Loop Integrals with SAILIR

    Apr 6, 2026David ShihScientific MLMemory-Efficient Inference

    NHETC, Dept. of Physics and Astronomy, Rutgers University, Piscataway, NJ 08854, USA

  3. Learning the S-matrix from data: Rediscovering gravity from gauge theory via symbolic regression

    Feb 16, 2026Nathan MoynihanHigh-Energy PhysicsSymbolic Regression

    Queen Mary, University of London, Mile End Road, London E1 4NS, UK