Generalizable Foundation Models for Calorimetry via Mixtures-of-Experts and Parameter Efficient Fine Tuning

This paper introduces a generalizable foundation model for calorimetry that leverages next-token transformer architectures combined with Mixture-of-Experts pre-training and parameter-efficient fine-tuning to enable modular, scalable, and computationally efficient simulation of particle showers across diverse materials and detector configurations without catastrophic forgetting.

The Gist

Imagine you are an architect trying to design the perfect building. To do this, you need to know exactly how the building will react to wind, rain, and earthquakes. In the world of particle physics, scientists are building "detectors" (giant cameras) to catch subatomic particles. To design these detectors, they need to run millions of computer simulations to see how particles crash into the detector's materials.

Traditionally, this simulation process is like trying to simulate every single raindrop hitting a roof using a supercomputer. It's incredibly accurate, but it takes so long and uses so much energy that it's becoming impossible to keep up with the demands of modern science.

This paper introduces a new, smarter way to do this using Artificial Intelligence, specifically a type of "Foundation Model" (a super-smart AI brain) designed for calorimetry (measuring particle energy). Here is how they did it, explained through simple analogies:

1. The Problem: The "One-Size-Fits-None" Dilemma

Imagine you have a master chef who is a genius at cooking steak. If you ask them to cook a fish, they might struggle because they've never done it. If you try to retrain them to cook fish, you might accidentally make them forget how to cook steak perfectly. This is called "catastrophic forgetting."

In physics, if you train a simulation AI on Tungsten (a heavy metal used in detectors), and then you want to simulate Lead, you usually have to start from scratch or risk ruining the Tungsten knowledge.

2. The Solution: The "Modular Kitchen" (Mixtures-of-Experts)

The authors built an AI that acts like a Master Chef with a team of specialized sous-chefs.

• The Base Model (The Master Chef): This is the core AI, trained on a "foundation" of physics. It knows the general rules of how energy moves and how particles interact. This part of the brain is frozen (locked in place) so it never forgets what it already knows.
• The Experts (The Sous-Chefs): Attached to the Master Chef are specialized modules called "Experts."
  • One expert specializes in Tungsten.
  • Another specializes in Tantalum.
  • A new one can be added for Lead.

When the AI needs to simulate a particle hitting Tungsten, it asks the "Tungsten Expert." When it needs to simulate Lead, it asks the "Lead Expert." The Master Chef stays the same; only the specific expert changes. This means you can add new materials without ever messing up the old ones.

3. Handling New Ingredients: "Low-Rank Adaptation" (LoRA)

What if you want to change the type of food being cooked? For example, switching from cooking Photons (light particles) to Electrons (charged particles). The rules of the kitchen change significantly.

Instead of firing the Master Chef and hiring a new one, the authors use a technique called LoRA (Low-Rank Adaptation). Think of this as giving the Master Chef a specialized apron and a new set of tools.

• The core brain (the chef's knowledge) stays the same.
• The apron (LoRA) adjusts how the chef thinks about the specific task (e.g., "Oh, electrons bounce differently than light").
• This is a tiny, lightweight adjustment that allows the AI to learn a new particle type quickly without needing to relearn everything from scratch.

4. The Result: Fast, Flexible, and Future-Proof

By combining these two tricks (Specialized Experts for materials + Specialized Aprons for particle types), the team created a system that is:

• Modular: You can add a new material or particle type by just plugging in a new "expert" or "apron."
• Efficient: It doesn't need to retrain the whole brain. It only learns the small new parts.
• Fast: They used tricks from the world of Large Language Models (like the ones powering chatbots) to make the AI run incredibly fast on graphics cards. It's now nearly as fast as older, simpler simulation methods but much more accurate.

The Big Picture

Think of this AI as a universal translator for particle physics.

• Old Way: You hire a new translator for every single language (material/particle) you encounter. It's expensive and slow.
• New Way: You have one brilliant translator who speaks the "universal language" of physics. When you need to translate a new dialect (a new material), you just hand them a small, specific dictionary (an Expert module). They instantly understand it without forgetting the previous languages.

This allows scientists to design better particle detectors faster, saving massive amounts of computing power and time, which is crucial for the next generation of experiments that will explore the deepest secrets of the universe.

Technical Summary

1. Problem Statement

Modern particle physics experiments face a critical bottleneck: the computational cost of high-fidelity Monte Carlo (MC) simulations (specifically Geant4) required for detector design and analysis is approaching the limits of available resources. As luminosity increases, the demand for simulation data outpaces computing capabilities.

• The Challenge: Traditional Deep Learning surrogates (GANs, VAEs, Diffusion models) often lack scalability and flexibility. They typically require retraining or separate models for every new detector material or particle type.
• The Risk: Standard "full fine-tuning" of a pre-trained model on new data leads to catastrophic forgetting, where the model loses its ability to accurately simulate the original configurations (e.g., a model fine-tuned for electrons may degrade its performance on photons).
• The Goal: Develop a Foundation Model (FM) for calorimetry that is:
  1. Generalizable: Capable of handling multiple materials and particle species within a single architecture.
  2. Modular: Allows for the addition of new capabilities (new materials/particles) without retraining the entire model.
  3. Efficient: Computationally competitive with existing fast simulation methods while maintaining high fidelity.

2. Methodology

The authors propose a Next-Token Transformer architecture inspired by Large Language Models (LLMs), adapted for physics simulation. The core innovation lies in combining Mixture-of-Experts (MoE) with Parameter-Efficient Fine-Tuning (PEFT).

A. Data Representation & Tokenization

• Input: Electromagnetic showers are voxelized into a $30 \times 30 \times 30$ grid.
• Tokenization:
  • Spatial: Each voxel coordinate $(x, y, z)$ is treated as a discrete token (vocabulary size $\approx 27k$).
  • Energy: Continuous energy values are discretized via linear binning into tokens (vocabulary size $\approx 25k$).
  • Sequence: Showers are generated sequentially, ordered by energy magnitude (descending), guided by the initial particle energy.
• Conditioning: The model is conditioned on the initial particle energy and a unique Particle Identifier (ID).

B. Model Architecture

The architecture consists of a frozen Pre-trained Backbone (Dual-sequence Transformer with Cross-Attention and Self-Attention blocks) and modular Adaptation Layers:

1. Rotary Positional Embeddings (RoPE): Used instead of learned positional embeddings to handle variable sequence lengths and relative positional dependencies efficiently.
2. Mixture-of-Experts (MoE) for Materials:
   • A Router directs inputs to specific "Expert" modules based on the material (e.g., Tungsten, Tantalum, Lead).
   • Strategy: The backbone remains frozen. New materials are added by introducing a single new expert and fine-tuning only that expert. This prevents interference with previously learned material distributions.
3. Parameter-Efficient Fine-Tuning (PEFT) for Particles:
   • Adapting to new particle species (e.g., switching from photons to electrons) requires more than just material adjustment because interaction dynamics change fundamentally.
   • LoRA (Low-Rank Adaptation): Applied to the Query, Key, Value, and Output projections within the attention blocks to modulate the structural relationships between tokens.
   • Modular Vocabularies: Separate output heads (vocabulary projections) are created for each particle species to handle particle-specific probability spaces without high-rank corrections to a shared matrix.

C. Training Strategy

• Pre-training: The backbone is trained on photons in Tungsten (W) and Tantalum (Ta).
• Adaptation (New Material): Freeze backbone and LoRA; add and train one new expert.
• Adaptation (New Particle): Freeze backbone; train LoRA modules and new vocabulary heads; add new material experts as needed.

3. Key Contributions

1. Generalizable Foundation Model: Successfully constructed a single model capable of generating high-fidelity showers for multiple materials (W, Ta, Pb) and particle types (photons, electrons) using a unified next-token prediction framework.
2. Modular Adaptation without Forgetting: Demonstrated that adding new materials via MoE and new particles via LoRA + Modular Vocabularies allows for additive knowledge integration. The base model parameters remain frozen, eliminating catastrophic forgetting.
3. Data Efficiency: Showed that the model can be fine-tuned to new materials (Lead) with extremely small datasets (as few as 1,000 samples) while maintaining high fidelity, crucial for scenarios where generating large MC datasets is expensive.
4. Computational Optimization: Proved that next-token transformer models can be computationally competitive with standard generative approaches (like Normalizing Flows) by applying LLM-specific inference optimizations (KV-caching, memory preallocation, CUDA graphs).

4. Results

• Fidelity: The model achieves distribution-level agreement with Geant4 ground truth for visible cell energy, total energy sum, hit multiplicity, and spatial shower profiles (longitudinal center of gravity, radial profile).
• Material Transfer:
  • The model successfully learned Tungsten and Tantalum simultaneously during pre-training.
  • Fine-tuning to Lead (Pb) using only 1k–10k samples yielded results comparable to models trained on full datasets, with no degradation in Tungsten/Tantalum performance.
• Particle Transfer:
  • Adapting from photons to electrons required ~50k samples for high fidelity.
  • The model captured the physical shift where electrons shower immediately (unlike photons which require conversion), though a minor post-hoc calibration was needed for longitudinal profiles in new materials due to low-rank constraints.
• Inference Speed:
  • Geant4: ~4100 ms (CPU).
  • Ours (Optimized Transformer): ~10.46 ms (Nvidia A100 GPU).
  • Speedup: The model is ~392x faster than Geant4 and competitive with other fast simulation methods (e.g., CaloClouds II, L2LFlows) while offering superior extensibility.

5. Significance

This work establishes Next-Token Transformer Architectures as a viable path for Physics-Aware Foundation Models.

• Detector Design Workflow: It enables a shift from repeated, CPU-intensive Geant4 simulations to a scalable workflow where a pre-trained backbone is incrementally adapted to new detector configurations (materials, geometries) and particle types using minimal GPU resources and data.
• Sustainability: By reducing the computational cost of simulation by orders of magnitude, it addresses the looming resource crisis in high-energy physics.
• Scalability: The modular design ensures that as new datasets become available, the model can expand its knowledge base additively without the risk of degrading previously learned physics, a critical requirement for long-term detector development and operation.

In summary, the paper bridges the gap between the flexibility of LLMs and the rigorous requirements of particle physics simulation, offering a robust, efficient, and extensible framework for future high-energy physics experiments.