Learning Scattering Amplitudes with Transformer Reinforcement Learning

Technical Summary: Learning Scattering Amplitudes with Transformer Reinforcement Learning

Problem Statement
The paper addresses the computational challenge of determining high loop-level scattering amplitudes in planar $\mathcal{N}=4$ Super Yang-Mills (SYM) theory. Traditional perturbative methods based on Feynman diagrams scale factorially with loop order and particle count, rendering them intractable for high orders. While recent work has framed the symbolic structure of these amplitudes as a sequence modeling problem solvable by Transformers, existing "Transformer-only" approaches suffer from two critical limitations:

1. Data Dependency: They require a vast majority of the final answer (e.g., 97% of coefficients for $L=6$) to be known a priori to serve as training data.
2. Consistency: Greedy sampling of the probability distribution often produces outputs that violate known physical relationships and symmetries, as the model predicts coefficients independently without enforcing global constraints.

The goal is to reconstruct the integer-valued coefficients of the symbol alphabet for the three-gluon form factor (specifically the $3g \to H$ amplitude) with significantly fewer known coefficients while guaranteeing that all physical constraints are satisfied.

Methodology
The authors propose a Transformer Reinforcement Learning (RL) algorithm that integrates exact linear relations and symmetries directly into the search process. The approach treats the reconstruction as a sequential search problem involving three distinct components:

1. Symbolic Representation and Constraints:

   • The amplitude is represented as a symbol $S[F^{(L)}]$ consisting of integer coefficients $C$ over sequences ("words") of length $2L$ drawn from a six-letter alphabet $\{a, b, c, d, e, f\}$.
   • The solution space is constrained by adjacency constraints (forbidden letter pairs and alternating structures) and linear relations (integrability conditions, causality, and all-loop relationships). These relations allow the deterministic inference of many coefficients from a partial assignment.

2. State Compression (Minimal Suffix Representation):

   • To handle the factorial growth of the state space, the authors employ a "minimal suffix representation." By analyzing relationships that act on word endings, they construct a compact basis of independent variables.
   • This reduces the state size by replacing suffixes with representative tokens, trading a larger token alphabet for a significantly shorter sequence length ($2L - K + 1$).

3. Algorithm Architecture:

   • Pretraining: A two-headed Transformer is pre-trained on a subset of known coefficients. The policy head learns to predict coefficients ($P(\text{coefficient}|\text{word})$), while the value head learns to estimate the remaining path length (via Mean Squared Error) to guide the search.
   • Reinforcement Learning Loop (MCTS): The algorithm operates in a loop:
     1. Selection: Identify a word with an unassigned coefficient that participates in the most relationships with only two unknowns.
     2. Proposal: The pre-trained Transformer proposes a distribution of candidate coefficients.
     3. Propagation: Exact linear relations are used to deterministically propagate the consequences of assigning a coefficient. This step resolves many other coefficients automatically.
     4. Search: When propagation reaches a fixed point with unresolved coefficients, Monte-Carlo Tree Search (MCTS) explores alternative assignments.
     5. Constraint Enforcement: Any assignment that violates a known relationship is treated as a "game over," pruning that branch of the search tree. This ensures every completed output is physically consistent.

Key Contributions

• Integration of Symmetries: Unlike previous Transformer-only methods, this algorithm incorporates derived symmetries and linear relationships as hard constraints within the learning loop, rather than relying solely on statistical learning.
• Hybrid Search Mechanism: The combination of Transformer-based coefficient proposal, deterministic propagation, and MCTS allows the system to navigate the combinatorial explosion of the state space.
• Data Efficiency: The method drastically reduces the fraction of the solution required as labeled pretraining data.
• Guaranteed Consistency: By treating violations as terminal states in the MCTS, the algorithm guarantees that every output satisfies the full set of imposed relationships, a feature not present in standard sequence modeling.

Results
The algorithm was tested on the $L=5$ symbol for the three-gluon form factor, which contains 12,543 words.

• Performance: The model successfully reconstructed the complete $L=5$ symbol using as little as 5% of the coefficients as known input.
• Comparison: This contrasts sharply with the Transformer-only approach, which required 97% of the symbols for training in the $L=6$ case.
• Efficiency: Propagation alone accounted for approximately 70% of the word assignments before MCTS intervention was necessary. The remaining work was handled by the learned priors of the Transformer.
• Verification: All generated solutions agreed with previously derived results (up to cyclic transformation) and satisfied every imposed relationship.

Significance and Claims
The paper claims that this approach is crucial for the generalization of machine learning to higher loop orders. Without the integration of exact relations and MCTS, the factorially scaling state sizes would make it impossible to compare results with those derived via other methods. The authors assert that their method allows for the derivation of high-loop results (specifically $L=5$) with a drastically smaller pretraining set while ensuring physical consistency.

The authors note a modest limitation: while their method uses significantly less computational power than recent results (specifically referencing Anthropic's $L=9$ results released shortly after their submission), their approach has not yet been demonstrated on loop 9. They state that extending the method to $L=9$ will be the subject of follow-up work.