\documentclass[nocopyright]{cimento}
\usepackage{graphicx}
\usepackage{subcaption}
\usepackage{float} 
\usepackage{amsmath}
%%%%%%%%%%%%%%%%
%
% VERY IMPORTANT 
%
% In order to set a Copyright owner please use and fulfill the following command 
%\setcopyright{CERN on behalf the XXXXX Collaboration}
%
%
%%%%%%%%%%%%%%%
\setcopyright{Copyright 2025 CERN for the benefit of the ATLAS Collaboration. CC-BY-4.0 license.}

\title{B-jet charge identification with graph neural networks for forward-backward asymmetry studies at ATLAS experiment\footnote{IFAE 2025 - “Poster” session}}

\shorttitle{B-jet charge identification with GNN\lowercase{s} for \(A^b_{FB}\) studies}
\author{Andrea Milici on the behalf of ATLAS collaboration}
\instlist{\item INFN, Sezione di Genova - Genova, Italy \\Universit\`a di Genova - Genova, Italy}
%% When only one author is present, please do not use the command \from{} near the author name.
\begin{document}

\maketitle

\begin{abstract}
The identification of b-flavor quarks that hadronize into jets, known as b-jet tagging, is essential for experiments at the LHC. The ATLAS \cite{bib:ATLAS} collaboration employs advanced machine learning algorithms based on the transformer architecture. The studies presented in this work aim to extend current algorithms, which already demonstrate excellent performance in terms of signal efficiency and background rejection, by incorporating the ability to distinguish between b and anti-b states, i.e., b-jet charge discrimination. \\
Preliminary results, obtained by extending existing transformer-based taggers, show a b/anti-b charge discrimination accuracy exceeding \(70\% \), a significant improvement over previous implementation in ATLAS. The development of such a powerful jet charge identification could provide an useful tool for a wide range of analysis. This is particularly true in the measurement of the b-quarks forward-backward asymmetry (\(A^b_{FB}\)) in final states with a Z boson produced in association with a b or anti-b quark, a process sensitive to the electroweak mixing angle of b quarks, first measured at LEP and still in tension with the Standard Model prediction.
\end{abstract}

\section{Introduction to Charge tagging}
\noindent \textit{Charge tagging} refers to the capacity of correctly identifying the electric charge of the parton (quark) that gives origin to a jet --- the complex and collimated pattern of tracks resulting from quark production and hadronization, as observed in the LHC detectors. 
Reconstructing this additional quantum number, \textit{i.e.} charge, allows interesting physics measurements such as the \textit{``B-quark forward-backward asymmetry"}, namely \(A^b_{FB}\). \\
First precisely measured at the LEP collider, this observable is sensitive to the electroweak couplings of the Z boson to \(b\)-quarks and it is directly related to the value of the weak mixing angle \(sin^2(\theta_W)\). The LEP measurement \cite{bib:LEP} showed a significant tension (around \( 2.5\sigma\)) with
the Standard Model prediction, suggesting a possible deviation in the \(Zb\bar{b}\) vertex — a result that remains intriguing and not fully explained. Recently, it has been shown that this asymmetry can also be studied at the LHC \cite{ref:Panizzo} looking at final states in \(Z+b\) associated production, with the Z boson decaying leptonically. To extract the asymmetry, one must reconstruct the angle between the lepton (or antilepton) from the Z decay and the \(\bar{b}\) (or b ) jet, in the Z boson rest frame, from this point comes the need of knowing the charge state of the b-jet. \\
Finally b-jet charge tagging can be used in analyses targeting di-b-jets resonances to suppress combinatorial background.

\section{Charge tagging methods}
\noindent Historically the main strategies for charge tagging, before the advent of machine learning methods, were \textbf{Jet Charge} and \textbf{Soft Lepton Tagging}, the former has been used extensively at LEP and the latter at CDF for measuring the charge of the top quark \cite{CDF:2010dwt}:
\begin{itemize}
    \item \textbf{Jet Charge} is calculated starting from  a variety of \(p_T\)-weighted sums of the reconstructed charge of  tracks \cite{Tokar:2017syr}:
\[
Q_k^{(1)} = \frac{1}{p_{\mathrm{T},J}^{\kappa}} \sum_{i \in \mathrm{Jet}} q_i \, (p_{\mathrm{T},i})^{\kappa},\quad
Q_k^{(2)} = \frac{\sum_{i \in \mathrm{Jet}} q_i \left| \vec{j} \cdot \vec{p}_i \right|^{\kappa}}{\sum_{i \in \mathrm{Jet}} \left| \vec{j} \cdot \vec{p}_i \right|^{\kappa}}, \quad
Q_k^{(3)} = \sum_{i \in \mathrm{Jet}} \left( \frac{E_i}{E_J} \right)^{\kappa} q_i,
\] 
where the index \textit{i} runs over all tracks associated to the jet, or, depending on the definition, only those matched to the leading hadron or to tracks from secondary or tertiary vertexes.
    \item \textbf{Soft Lepton Tagging} (SLT) relies on the direct link between the charges of b-quark and lepton from semileptonic decays of B-hadrons.
\end{itemize}
Jet Charge can be applied to all jets, providing high statistics but relatively low purity (\(<60\%\)). On the contrary, Soft Lepton Tagging (SLT) method offers higher purity, limited mainly by fake tags arising from cascade decays (\(b\rightarrow c\rightarrow \ell \nu X\)) --- which can be in principle kinematically suppressed. However, SLT is statistically bounded  by the semileptonic branching fraction \(BR(b\rightarrow \ell\nu X) \approx 10\%\). \\
Furthermore, in 2015, ATLAS explored multivariate approaches to charge tagging with \textbf{Jet Vertex Charge} (JVC) algorithm \cite{ATLAS:2015jhb} \cite{ATLAS:2018lhe}. This method employed a Multi-Layer Perceptron (MLP) that combined various jet charge observables and low-level inputs to produce a score for identifying the b-jets' charge. The performance achieved by the JVC algorithm reached an efficiency of approximately \(63\%\) in correctly identifying the electric charge of the originating quark.

\section{GN2 with Charge classification}
\noindent The introduction of transformer architecture is revolutionizing flavour tagging \cite{bib:GN2}; these algorithms enable a significantly improved background rejection compared to previous taggers (DL1). Since the trend is to provide increasingly low-level inputs to the network, it becomes challenging to understand what the model is actually learning about the underlying physics and the jet substructure. The use of auxiliary tasks in combination with the main one allows the extraction of physical information as well as improvements in overall convergence and performance.
In this work we introduce an auxiliary task for b-jet charge identification in the GN2 tagger, alongside the existing ones for Track Origin prediction and Vertexing, as summarized in figure \ref{GN2}. 

\begin{figure}[H]   
    \centering
    \includegraphics[width=0.9\linewidth]{images/GN2.eps}
    \caption{\centering GN2 architecture \cite{bib:GN2}}
    \label{GN2}
\end{figure}

\noindent Since the interest here is to tag the charge of the quark at the parton level, particular attention must be paid to the truth labels used to train the task. These labels can be categorized into two types: parton-based and hadron-based truth variables. We chose to focus on the latter ones, as parton-based definitions are known to be not infrared and collinear (IRC) safe, and to ensure consistency with the truth labelling scheme adopted by the ATLAS Flavour Tagging (FTAG) group.
\noindent Within the hadron-based category, a further distinction can be made between variables associated with the initial hadrons (i.e., before any possible meson oscillation) and those corresponding to the final stable hadrons. This distinction is particularly relevant for charge tagging, as neutral B-mesons can undergo flavour oscillations that alter the apparent charge of the jet spoiling the performance evaluation. The quark charge is ultimately inferred from the quark content of the b-hadron labelling the jet before possible oscillations.

\section{Performance}
\noindent The current implementation of the algorithm outputs three probabilities: \( p_b,p_{\bar{b}}, p_{other}\) representing the probability that a jet originates from a \(b\)-quark, a \(\bar{b}\)-quark or everything else. Based on these values, we define a log-likelihood discriminant to construct a continuous score as in Eq.~\ref{discr}:
\begin{equation}
D^b_{\text{charge}} = \log\left(\frac{p_b}{(1 - f_o)p_{\bar{b}} + f_o p_{other}}\right)
\label{discr}
\end{equation}
Where we take \(f_o=0\).
\begin{figure}[h!]
    \centering
    \begin{subfigure}[t]{0.48\linewidth}
        \centering
        \includegraphics[width=\linewidth]{images/ChargeID_ROC_pre_btagging_selection.eps}
        \caption{ROC curve for the charge identification of jets originating from b-quarks.}
        \label{fig:ChargeROC}
    \end{subfigure}
    \hfill
    \begin{subfigure}[t]{0.48\linewidth}
        \centering
        \includegraphics[width=\linewidth]{images/png2pdf.eps}
        \caption{ROC curves for flavour tagging.} % Sostituisci con una descrizione reale
        \label{fig:ftagROC}
    \end{subfigure}
    \caption{Charge tagging and Flavour tagging performance.}
    \label{fig:ROC}
\end{figure}

\noindent In Figure \ref{fig:ChargeROC}, ROC curves are shown for different training configurations. \textit{InitialHadrons} and \textit{FinalHadrons} refer to the type of truth labels used during training. Furthermore, the GN2 BSplitted setup investigates the possibility of identifying the b-jet charge directly within the main classification task by splitting the b-jet class. \\
\noindent The performance achieved in terms of efficiency of correctly identifying the charge of a b-jet ranges from \(70\%\) to \(72.5\%\) depending on the b-tagging working point. The quoted efficiency is calculated starting from the cut on \(D^b_{charge}\) (Eq.~\ref{discr}) that guarantees equal discrimination both for \(b\) and \(\bar{b}\), showing significant improvements compared to the previous tagger. Figure~\ref{fig:ftagROC} shows that the flavour tagging performance remains stable, with no significant degradation due to the introduction of the b-charge auxiliary task. 
\begin{figure}[h!]
    \centering
    % Prima riga
    \begin{subfigure}[b]{0.45\textwidth}
        \centering
        \includegraphics[width=\linewidth]{images/hadrons1.eps}
        \caption{\(B^+-B^-\)}
        \label{fig:sub1}
    \end{subfigure}
    \hfill
    \begin{subfigure}[b]{0.45\textwidth}
        \centering
        \includegraphics[width=\linewidth]{images/hadrons2.eps}
        \caption{\(B^0-\bar{B}^0\)}
        \label{fig:sub2}
    \end{subfigure}
    % Seconda riga
    \begin{subfigure}[b]{0.45\textwidth}
        \centering
        \includegraphics[width=\linewidth]{images/hadrons3.eps}
        \caption{\(\Xi^+_b-\Xi^-_b\)}
        \label{fig:sub3}
    \end{subfigure}
    \hfill
    \begin{subfigure}[b]{0.45\textwidth}
        \centering
        \includegraphics[width=\linewidth]{images/hadrons4.eps}
        \caption{\(B_s^0-\bar{B}_s^0\)}
        \label{fig:sub4}
    \end{subfigure}

    \caption{Discriminant \(D^b_{charge}\) on different hadron species.}
    \label{fig:bhadrons}
\end{figure}
\noindent Figure \ref{fig:bhadrons} shows the discriminant \(D^b_{charge}\) (Eq.~\ref{discr}) calculated for different types of b-hadrons. A good separation is observed for charged b-hadrons; the worse separation for the \(B^0_s\) compared to the \(B^0\) reflects its faster oscillation frequency.
\begin{figure}[h!]
    \centering
    \begin{subfigure}[t]{0.48\linewidth}
        \centering
        \includegraphics[width=\linewidth]{images/discriminant.eps}
        \caption{Discriminant \(D^b_{charge}\).}
        \label{fig:Discriminant}
    \end{subfigure}
    \hfill
    \begin{subfigure}[t]{0.48\linewidth}
        \centering
        \includegraphics[width=\linewidth]{images/histogram_bjets_probability.eps}
        \caption{Output probabilities distribution.} 
        \label{fig:Probs}
    \end{subfigure}
    \caption{Inclusive discriminant and probabilities distribution.}
    \label{fig:ROC}
\end{figure}

\noindent Finally, figures \ref{fig:Discriminant} and   \ref{fig:Probs} show the inclusive discriminant and the distribution of output probabilities.
\section{Conclusions}
\noindent This work presents a novel approach to charge tagging using state-of-the-art deep learning architectures. Leveraging the attention mechanism, the network captures complex correlations among tracks — this is particularly important for jet charge reconstruction, where hadronization processes play a crucial role \cite{Waalewijn:2012sv}. The final aim of the work will be to provide a reliable and calibrated tool that can be effectively used in analyses.
\newpage

\acknowledgments
\begin{thebibliography}{0}

\bibitem{bib:ATLAS}
ATLAS Collaboration,  
``The ATLAS Experiment at the CERN Large Hadron Collider,''  
JINST \textbf{3}, S08003 (2008),  
doi:10.1088/1748-0221/3/08/S08003.

\bibitem{bib:LEP}
The ALEPH Collaboration, the DELPHI Collaboration, the L3 Collaboration, the OPAL Collaboration, the SLD Collaboration, the LEP Electroweak Working Group, the SLD electroweak, heavy flavour groups,
``Precision Electroweak Measurements on the Z Resonance``, Phys.Rept.427:257-454,2006

\bibitem{ref:Panizzo}
G.~Panizzo,
``Testing new physics with bottom quarks at LHC: a pragmatic approach.''  

%\bibitem{ATLAS:2020lbz}
%ATLAS Collaboration,
%``Measurement of the $CP$-violating phase $\phi_s$ in %$B^0_s \to J/\psi\phi$ decays in ATLAS at 13 TeV,''  
%Eur.\ Phys.\ J.\ C \textbf{81}, 342 (2021),  
%doi:10.1140/epjc/s10052-021-09011-0  
%[arXiv:2001.07115 [hep-ex]].

\bibitem{CDF:2010dwt}
CDF Collaboration,``Exclusion of an Exotic Top Quark with -4/3 Electric Charge Using Soft Lepton Tagging,''  
Phys.\ Rev.\ Lett.\ \textbf{105}, 101801 (2010),  
doi:10.1103/PhysRevLett.105.101801,arXiv:1006.4597.

\bibitem{Tokar:2017syr}
S.~Tokar on the behalf of ATLAS and CMS Collaborations,  
``Jet charge determination at the LHC,''  
ATL-PHYS-PROC-2017-017.

\bibitem{ATLAS:2015jhb}
ATLAS Collaboration,  
``A new tagger for the charge identification of b-jets,''  ATL-PHYS-PUB-2015-040.

\bibitem{ATLAS:2018lhe}
ATLAS Collaboration,  
``Measurement of the Jet Vertex Charge algorithm performance for identified $b$-jets in $t\bar{t}$ events in $pp$ collisions with the ATLAS detector,''  
ATLAS-CONF-2018-022.

\bibitem{bib:GN2}
ATLAS Collaboration, ``Transforming jet flavour tagging at ATLAS'',arXiv:2505.19689.

\bibitem{Waalewijn:2012sv}
W.~J.~Waalewijn,  
``Calculating the Charge of a Jet,''  
Phys.\ Rev.\ D \textbf{86}, 094030 (2012),  
doi:10.1103/PhysRevD.86.094030  
[arXiv:1209.3019 [hep-ph]].


\end{thebibliography}
\end{document}
