Transformers Learn Physical Symmetries in Cosmic Ray Data

By visualizing the internal mechanisms of the networks, the team shows that Transformers do not merely act as black boxes but can encode specific physical symmetries and source correlations directly into their learned parameters.

The first scenario focuses on the reconstruction of air shower properties using a hexagonal sensor array. Ultra-high-energy cosmic rays, likely composed of ionized atomic nuclei ranging from protons to iron, collide with atoms in the Earth’s atmosphere to generate air showers containing billions of secondary particles. These showers are detected by ground-level sensors, typically arranged in a hexagonal pattern. The study examines whether the Transformer architecture, which lacks explicit symmetry constraints in its design, can learn the azimuthal rotational symmetry inherent in these air showers. The network was trained on approximately 500,000 simulated air showers resulting from primary cosmic nuclei with mass numbers A=1, 4, 16, and 56.

To investigate this, the researchers analyzed the trainable positional encodings within the Vision Transformer architecture. Positional encodings are vectors added to the input data to provide information about the arrangement of the sensors. In this study, the two-dimensional structure of the hexagonal sensor array is flattened into a one-dimensional vector, requiring the network to re-learn the spatial relationships through these encodings. The team calculated the normalized scalar product, similar to the cosine of the angle, between the positional encoding vectors of different sensors. This metric reveals the similarity of the encodings across the array.

The results showed that the Transformer successfully learned the rotational symmetry of the hexagonal sensor structure. This pattern held true regardless of which sensor was used as the reference point, indicating that the network had internalized the azimuthal symmetry of the air showers purely from the simulated data. The study notes that while the network architecture and loss function did not contain explicit information about symmetries, the learned positional encodings effectively captured this domain-specific physical property to improve reconstruction accuracy.

The second scenario addresses the challenge of identifying the galactic sources of ultra-high-energy cosmic rays. Astronomical observations indicate that the Milky Way’s magnetic field significantly deflects ionized atomic nuclei, making it difficult to trace particles back to their origins based on arrival direction alone. The study uses a Transformer network as a pre-processing tool for a conditional invertible neural network (cINN). The goal is to determine which cosmic particles originate from a catalog of gamma-ray active galactic nuclei (gamma-AGNs) and which are background particles, while simultaneously adjusting the deflection angles and strengths of the galactic magnetic field model.

The network architecture includes a fully connected embedding layer, a Transformer with eight heads, and a cINN. The Transformer processes the observed cosmic particles, characterized by zenith, azimuth, energy, and shower depth, encoded in a 128-dimensional latent space. The team trained the network using on the order of 10^6 astrophysical simulations, each containing approximately 4,000 cosmic particles. In these simulations, 10% of the particles were designated as signal particles from the galaxy catalog, while 90% were background particles. The signal particles were deflected differently in each simulation through variations of harmonic coefficients to train the invertible network.

To understand how the Transformer distinguishes signal from background, the researchers visualized the attention values assigned to the cosmic particles. Attention mechanisms allow the network to evaluate correlations between input information, effectively highlighting which parts of the input are most relevant for the task. By summing the attention values and mapping them to a sky map using the healpix coordinate system, the team created visual representations of where the Transformer focused its “attention.” The analysis revealed that each of the eight Transformer heads concentrated on a specific region in the sky, corresponding to the expected arrival directions of signal particles from the catalog galaxies.

The study further quantified the discrimination power of the Transformer by comparing the aggregated attention values for signal particles versus background particles in previously unknown astrophysical scenarios. The results showed that the attention values of signal cosmic rays were well separated from those of background cosmic rays. This separation was consistent across multiple Transformer heads, indicating that the network had learned to assign higher attention to particles originating from the given galaxy catalog. The high attention regions in the sky maps aligned with the coherent deflections predicted by the galactic magnetic field model, modified by the harmonic coefficients inferred by the cINN.

An integrated gradients approach was used to examine the importance of the input variables in the attention mechanism. This method calculates the gradient of the output with respect to the input across the complete network structure, providing a measure of feature importance. The findings indicated that directional information dominates the attention mechanism. After normalizing the total scalar importance to unity, energy accounted for approximately 5% of the attention weight, while shower depth contributed about 2%. Both estimates carried an uncertainty of roughly 40%. This suggests that the Transformer primarily relies on the spatial distribution of particles to identify potential sources, with energy and mass-related information playing a secondary role in the attention process.

The study highlights the capability of Transformer networks to learn complex physical structures from simulation data without explicit architectural constraints. In the air shower reconstruction scenario, the network encoded azimuthal symmetry into its positional encodings, demonstrating an ability to capture rotational invariance. In the source identification scenario, the attention mechanism focused on specific sky regions corresponding to galactic sources, effectively separating signal from background. These findings provide deeper insight into the internal workings of Transformers in scientific applications, moving beyond global sensitivity analyses to targeted visualizations of learned features.

The study acknowledges that while the Transformer architecture does not include explicit symmetry information, the training process allows it to encode such domain-specific knowledge.

By examining positional encodings and attention maps, researchers can gain a more intuitive understanding of how the network processes data.

The publication, titled “What Exactly Did the Transformer Learn from Our Physics Data?”, was released on September 15, 2025, in Volume 9 of Computing and Software for Big Science. It is an open-access article that provides a detailed analysis of two Transformer applications in ultra-high-energy cosmic ray physics. The work underscores the potential of Transformer networks to not only achieve high performance in scientific tasks but also to learn and encode fundamental physical principles, such as symmetry and source correlation, from data alone.

Subject of Research: What Exactly Did the Transformer Learn from Our Physics Data?

Article Title: What Exactly Did the Transformer Learn from Our Physics Data?

Article References: Erdmann, M., Langner, N., Schulte, J., & Wirtz, D. (2025). What Exactly Did the Transformer Learn from Our Physics Data?. Computing and Software for Big Science, 9(1), Article 16. https://doi.org/10.1007/s41781-025-00145-4

Image Credits: AI Generated

DOI: 10.1007/s41781-025-00145-4

Keywords: Exactly, Transformer, Learn, Physics, Data, scientific research