BRICKS: Compositional Neural Markov Kernels for Zero-Shot Radiation-Matter Simulation

This paper introduces BRICKS, a differentiable, compositional neural surrogate based on hybrid discrete-continuous transformers and Riemannian Flow Matching that enables zero-shot, high-speed simulation of radiation-matter interactions by composing next-particle prediction kernels to model unseen large-scale material distributions.

The Gist

Imagine you are trying to predict what happens when a single billiard ball hits a complex, multi-layered wall made of different materials. In the real world of physics, this is incredibly hard to calculate because the ball might bounce, shatter into smaller pieces, create heat, or trigger a chain reaction of other tiny particles.

Traditionally, scientists use "mechanistic simulators" to solve this. Think of these simulators as a super-detailed, slow-motion camera that tracks every single tiny collision, one by one, for every single particle. It's accurate, but it's like trying to count every grain of sand on a beach to understand the shape of the dunes. It takes a massive amount of computer power and time.

The paper introduces BRICKS, a new way to do this simulation that is faster, smarter, and more flexible. Here is how it works, broken down into simple concepts:

1. The "Lego" Philosophy (Composition)

The core idea of BRICKS is composition. Imagine you have a small box of Lego bricks. If you understand exactly how one specific brick snaps onto another, you don't need to be shown a picture of every possible castle, spaceship, or house to know how to build them. You just need to know the rule for connecting the bricks.

• Old Way: Train a computer to recognize a picture of a specific, finished castle (a specific material setup). If you want to simulate a different castle, you have to retrain the computer.
• BRICKS Way: Train the computer on the "rule" of how one particle interacts with a small chunk of material (a "kernel"). Once it learns this rule, it can snap these rules together to simulate any new material shape it has never seen before. This is called Zero-Shot Generalization—it works on new things without needing extra practice.

2. The "Next-Particle" Predictor

Instead of simulating the whole journey of a particle through a massive wall, BRICKS acts like a predictive engine for the next step.

• You give it: "Here is a particle coming in, and here is the material it's hitting."
• It answers: "Here is the new set of particles that come out, and here is the energy left behind in the material."

It treats the interaction like a story where you only need to know the current scene to predict the next scene, rather than writing the whole book at once.

3. The "Hybrid Brain" (The Model)

To make these predictions, the team built a special AI brain using Transformers (the same technology behind modern chatbots). However, this brain is unique because it handles two types of information at once:

• Discrete (The "What"): It counts how many new particles are created (e.g., "I see 2 electrons and 1 photon"). This is like counting apples.
• Continuous (The "How"): It predicts the exact speed, direction, and energy of those particles. This is like measuring the weight of the apples.

The paper uses a technique called Riemannian Flow Matching. Think of this as a smooth, mathematical river that guides the AI from a state of "random noise" to a state of "accurate prediction." It ensures the AI doesn't just guess; it learns the precise probability of every outcome, allowing it to be "differentiable" (meaning scientists can use the math behind the prediction to optimize other things later).

4. The "CaloBricks" Dataset

To teach this AI, the researchers couldn't just use old data. They needed a new kind of textbook. They created CaloBricks, a massive dataset of 20 million simulated interactions.

• They shot electrons, positrons, and photons at cubes of Argon gas (a common material in physics detectors) with varying densities.
• They recorded exactly what went in and what came out.
• This dataset is now being released to help other scientists train similar models.

5. The Results: Speed and Stability

The team tested BRICKS in two ways:

• Single Step: When looking at just one interaction, the AI's predictions were almost identical to the slow, traditional simulators.
• Chained Steps: They let the AI run the simulation over and over (like a chain reaction). Even after many steps, the errors didn't pile up and ruin the result. It remained stable.

The Big Win:
The most exciting result is speed. Because the AI runs on specialized computer chips (GPUs) and skips the need to simulate every tiny micro-collision, it is significantly faster than the traditional CPU-based methods, especially when dealing with dense materials where the old method would have to do millions of calculations.

Summary

BRICKS is like teaching a computer the "grammar" of particle physics rather than memorizing every "sentence" (simulation). By learning the basic rules of how particles interact with small chunks of matter, the model can instantly compose those rules to simulate complex, unseen environments. It offers a faster, more flexible, and mathematically transparent way to simulate radiation, which is crucial for fields like particle physics, nuclear engineering, and medical physics.

Technical Summary

Technical Summary: BRICKS – Compositional Neural Markov Kernels for Zero-Shot Radiation-Matter Simulation

Problem Statement

Simulating radiation-matter interactions is a critical task across particle physics, nuclear engineering, space exploration, and medical physics. Current state-of-the-art approaches rely on mechanistic simulators (e.g., Geant4), which model these interactions as iterative stochastic processes. These simulators track particles through a "basic vocabulary" of quantum interactions, strictly adhering to the Markovian nature of the physics (where the next state depends only on the current particle state and local material properties).

However, mechanistic simulators face significant limitations:

1. Computational Cost: They require a vast number of iterations to simulate large-scale geometries, consuming substantial CPU resources (e.g., projected to use 20–30% of LHC CPU budgets).
2. Lack of Differentiability: They function as pure samplers, lacking access to tractable likelihoods and gradients, which limits their utility in downstream tasks like parameter inference or optimization.
3. Rigidity: Existing neural surrogate models typically learn to emulate full-scale detector responses for fixed geometries. Consequently, they lose the compositional nature of the underlying physics, failing to generalize to unseen material distributions without retraining.

The goal is to create a surrogate model that retains the compositional, zero-shot generalization capabilities of mechanistic simulators while offering the speed, differentiability, and likelihood access of neural networks.

Methodology

The authors propose BRICKS (Broadly Reusable InteraCtion Kernel Surrogates), a strategy that distills radiation-matter physics into a composable neural Markov kernel.

1. Probabilistic Model Formulation

Instead of simulating full geometries, the model learns a local transition kernel $p(\{z\}_{out}, \delta | z_{in}, \lambda)$, where:

• $z_{in}$: The incident particle (type, position, momentum, energy).
• $\lambda$: Local material properties (e.g., density).
• $\{z\}_{out}$: A variable-sized set of outgoing particles.
• $\delta$: A material side effect (e.g., energy deposition).

The simulation of a large geometry is achieved by autoregressively composing this kernel: the output particles of one step become the input for the next, traversing the material volume. This approach exploits the strict Markovian nature of particle transport, allowing the model to generalize to any geometry composed of the learned building blocks without retraining.

2. Neural Architecture

The kernel is implemented as a hybrid discrete-continuous conditional generative model using two distinct transformer-based components:

• Cardinality Model (Discrete): Predicts the number of outgoing particles for each type (e.g., $n_{e^-}, n_{e^+}, n_{\gamma}$).

  • Uses a causal autoregressive transformer.
  • Predicts particle-type counts in a fixed order rather than a variable-length sequence, reducing inference costs.
  • Inputs are encoded via embedding tokens for particle type and continuous material features.

• Continuous Feature Model (Continuous): Predicts the specific properties of the generated particles (position, momentum) and the energy deposition $\delta$.

  • Uses Riemannian Conditional Flow Matching (CFM) with a transformer backbone.
  • Manifold Handling: The model accounts for the non-trivial product manifold of the data:
    • Material response $\delta \in [\delta_{min}, \delta_{max}] \subset \mathbb{R}$.
    • Particle position on a cube surface (homeomorphic to $S^2$).
    • Particle momentum in $\mathbb{R}^3$.
  • Likelihood & Gradients: By estimating an instantaneous velocity field and integrating via an ODE solver, the model provides tractable likelihoods and gradients, a key requirement for future downstream applications.
  • Permutation Invariance: The architecture ensures equivariance to permutations of particles of the same type while distinguishing between different particle types via role-dependent token embeddings.

3. Dataset: CaloBricks

To train this kernel, the authors introduced CaloBricks, a novel dataset containing 20 million simulated events.

• Setup: Particles (electrons, positrons, photons) interact with a 10 cm Argon cube.
• Variability: Random incident positions, momentum directions, energies (20–300 MeV), and material densities (0.5–10 g/cm³).
• Generation: Created using the Geant4 toolkit.
• Scope: Includes incident conditions, outgoing particle data, and energy deposition. The dataset is designed to be expandable to other materials (e.g., Copper, Lead).

Key Contributions

1. Composable Probabilistic Model: Development of a multi-factor probabilistic model for particle-matter interactions that functions as a Markov kernel, enabling zero-shot generalization to unseen material distributions.
2. CaloBricks Dataset: Release of a large-scale (20M events) dataset of particle-material interactions with detailed outgoing particle data, filling a gap in existing benchmarks that typically only provide aggregate energy deposition.
3. BRICKS Architecture: Implementation and training of a mixed discrete-continuous generative model using Riemannian Flow Matching. The model is evaluated both as a standalone predictor and as a kernel under autoregressive rollouts.

Results and Evaluation

The model was evaluated against the ground-truth Geant4 simulator (G4) on single-kernel execution and multi-step autoregressive rollouts.

Single-Kernel Performance

• Distributional Fidelity: The model demonstrates strong agreement with Geant4 across marginal feature distributions and two-feature correlations.
• Metrics: Evaluated using Maximum Mean Discrepancy (MMD), Energy Distance (ED), and a trained neural classifier (AUC).
  • The classifier AUC scores are close to 0.5 (random guessing), indicating the model's distribution is indistinguishable from the ground truth.
  • MMD and ED scores are significantly lower than those obtained when comparing Geant4 to naive base distributions.
• Zero-Shot Generalization: The model maintains high fidelity across various conditioning sub-manifolds (varying angles, positions, energies, and densities) not explicitly seen during training in specific combinations.

Inference Speed

• Acceleration: The model achieves a significant speed-up on GPU hardware compared to CPU-bound Geant4 simulations.
• Scalability: The performance gap widens as material density increases (which increases the number of microscopic transitions required in mechanistic simulation).
• Trade-offs: The authors analyzed the Pareto frontier between inference time and quality by varying the ODE solver (Euler, Midpoint, RK4) and integration steps.

Autoregressive Stability (Zero-Shot Composition)

• Multi-Rollout: The model was tested in synthetic scenarios with varying material densities over multiple rounds (e.g., High-Low, Alternating, Random).
• Stability: The model successfully simulated the evolution of energy deposition and particle multiplicity over 10 rounds without catastrophic error accumulation.
• Metric: The MMD score for the concatenated rollout summary vectors remained low, comparable to the statistical floor of Geant4 comparing itself to itself.

Significance and Claims

The paper claims that BRICKS represents a significant step toward general-purpose neural surrogates for particle-matter simulation. Its primary significance lies in:

• Preserving Compositionality: Unlike previous surrogates tied to fixed geometries, BRICKS learns the local interaction rule, allowing it to simulate arbitrary large-scale geometries composed of known building blocks in a zero-shot manner.
• Differentiability and Likelihoods: By using Flow Matching, the model provides access to likelihoods and gradients, opening avenues for applications (e.g., inverse problems, optimization) that are intractable with standard samplers.
• Feasibility: The work demonstrates that it is possible to distill complex medium-scale physics into an amortized neural representation that is both accurate and faster than mechanistic simulation on a per-event basis.

The authors remain modest, noting that while the model is sound and shows promising autoregressive stability, it is not yet a full production-grade replacement for mechanistic simulators. Future work is required to develop a full 3D autoregressive framework, extend the kernel to more complex material building blocks, and optimize inference speed further (e.g., via KV caching strategies) for production deployment.