Adversarial Training for Mass Decorrelation

Adversarial Training for Mass Decorrelation

If you have trained a jet tagger or di-jet classifier and then looked at the invariant mass distribution in your background sample after applying a cut on the classifier score, you have probably seen this: a bump or dip that was not there before. The classifier learned that signal events cluster in a particular mass window, so it carved a corresponding shape into the background. In a bump hunt, that sculpted feature is indistinguishable from a signal. Adversarial training decorrelation is one of the most practical tools for preventing this, and this article walks you through how it works and how to set it up.

Why Classifiers Sculpt Mass

A standard classifier optimizes one objective: separate signal from background as cleanly as possible. Invariant mass is often one of the most powerful discriminating variables, directly or indirectly — through correlated observables such as jet mass, substructure variables, or even track multiplicity. The classifier is not told to ignore mass; it simply uses every correlation available to it.

The result is a score that depends on mass. When you apply a score threshold, you implicitly apply a mass-dependent selection. The mass distribution of surviving background events is no longer flat (or smoothly falling, or whatever your prior was). For a bump hunt, where the background model is derived by fitting the sidebands, this sculpted shape can produce a fake local excess.

Mass decorrelation machine learning techniques break this correlation during training, rather than trying to correct for it afterward.

The Adversarial Idea

Adversarial training borrows from generative adversarial networks but serves a different purpose here. Instead of generating realistic samples, you attach a second network — called the adversary or discriminator — whose job is to predict the mass of an event from the classifier's output score alone. The two networks train against each other:

  • The classifier tries to maximize signal–background separation.
  • The adversary tries to predict mass from the classifier score.
  • The classifier is simultaneously penalized for making the adversary's job easy.

In physics language: you are asking the classifier to be powerful (high efficiency for signal, low efficiency for background at fixed score threshold) while also being mass-agnostic (the score distribution of the background should look the same regardless of which mass bin you are in).

This is formalized by adding the adversary's loss to the classifier's loss with a negative sign and a tunable weight λ. Minimizing the total loss means maximizing signal separation while minimizing how much mass information leaks through the score.

Setting It Up in Practice

Step 1: Preprocess and define your mass variable

Choose the variable you want to decorrelate against — typically the invariant mass of the jet or dijet system, or the reconstructed candidate mass. Bin it into a moderate number of regions covering your search window. The adversary will try to classify which bin an event falls into.

Step 2: Build the classifier network

Use whatever architecture fits your problem — a dense network on high-level features, a graph network on particle-flow candidates, or anything else. The classifier outputs a single score between 0 and 1. For a broader introduction to network architectures used in HEP, the full HEP ML course covers these building blocks systematically.

Step 3: Build the adversary network

The adversary takes the classifier score as its only input and outputs a probability distribution over your mass bins (a softmax output with one node per bin). It should be shallow — two or three dense layers is typical. A powerful adversary gives the classifier a harder constraint, which is what you want.

Step 4: Define the combined loss

Total loss = classifier loss − λ × adversary loss

The minus sign is the key: the classifier is penalized when the adversary succeeds. λ controls the trade-off between discrimination power and decorrelation quality. Larger λ forces stronger decorrelation at the cost of some classification performance. You will need to scan λ and evaluate both goals.

Step 5: Alternate or jointly train

The standard approach alternates gradient steps. First update the adversary weights with the classifier frozen; then update the classifier weights (with the combined loss) with the adversary frozen. Repeat. Some implementations train jointly in a single step using gradient reversal layers, which flip the sign of the gradient flowing from the adversary into the classifier automatically.

Step 6: Validate with the Jensen–Shannon divergence or a χ² test

After training, slice your background sample into mass bins and overlay the score distributions. If decorrelation is working, those distributions should be nearly identical across bins. A χ² comparison of mass distributions before and after score cuts is the physics-intuitive check: the cut should not change the mass shape.

Common Pitfalls

λ too large: The classifier loses most of its discrimination power. λ too small: Sculpting is reduced but not eliminated. Adversary too weak: The classifier can fool it without truly becoming mass-agnostic — use a fresh, well-trained adversary at each evaluation. If you are also working on uncertainty-aware modeling, our discussions of systematic uncertainties in ML are directly relevant here.

One-Line Takeaway

Adversarial training decorrelation makes your classifier powerful and mass-blind at the same time — by explicitly punishing the network whenever its score leaks information about mass, you get a tagger you can trust in a bump hunt.

Want to go deeper?

Machine Learning for High Energy Physics: The Complete Course takes you from first principles to a defensible result in 6 structured modules. $97, 30-day guarantee.

See the course →

Not ready yet? Grab Module 1 free →