Emergent Explainability: Adding a causal chain to neural network inference

Explainable artificial intelligence (xAI) remains one of the foremost challenges in the field, crucial to equitable and fair application, particularly in sensitive domains like healthcare. For examples, understanding the impact of race in AI models is a common problemHuang et al. [2022]. Mitigating potential bias that exists within data and its effect on treatment begins with understanding the AI model. The core issue lies in the inherent opacity of advanced AI models, especially deep learning algorithms, which operate as "black boxes". These complex models can make highly accurate predictions or decisions, but understanding the causal chain leading to their output is not straightforward. This lack of transparency raises concerns about trustworthiness, accountability, and ethics, especially when AI is employed in critical areas such as medical diagnostics, treatment planning, or patient care management. This work addresses concerns of xAI with the goal of enabling AI to be a vital and interpretable tool for clinicians and patients alike, fostering trust and enabling informed decision-making. The novel framework also has many potential applications beyond healthcare, with the potential to reshape the way we interact with AI.

In the healthcare sector, xAI methods face significant challenges. To address concerns of safety, equity and trustworthiness, an understanding of causality is crucial. Simpler models like decision trees trade complexity for clarity, essential where decision rationales are crucial. However, complex models using tools like LIME Ribeiro et al. [2016] and SHAP Lundberg and Lee [2017] often yield limited, sometimes misleading insights. Visualization techniques in deep learning, such as attention maps, provide helpful interpretations but lack precise correlations with neural network processes Erhan et al. [2009], Simonyan et al. [2013]. Feature importance measures, while identifying key decision factors, can suffer from bias and lack causal clarity Zhou et al. [2016], Selvaraju et al. [2017]. These methods, enhancing reliability and aiding compliance, struggle with a fundamental trade-off between accuracy and explainability. This leads to technical complexities and the risk of oversimplified interpretations, underscoring the need for more refined and transparent explainability techniques in high-stakes healthcare AI applications.

This work focuses on enhancing xAI in healthcare, a key area in need of growth and equitable application. It aims to develop a novel and practical framework for AI explainability, particularly in healthcare diagnostics and patient care. The proposed research will imbue the AI model with a causal communication channel that can be used for validation and transparency of output. The expected outcomes include improved AI clarity and trust between clinicians and patients, with the potential to inform policy in AI ethics. This effort is geared towards contributing meaningful research and reshaping the way we use AI as a tool.

Emergent communication (EmCom) in AI Brandizzi [2023], a relatively nascent field, presents unique challenges and opportunities in the context of explainability. The principle behind EmCom is to have two or more agents exchanging messages in order to solve a task, see Fig.1. Altering the task set-up can alter the compositionality and generalisation of the generated language Mu and Goodman [2021], Lowe et al. [2020]. Current research focuses on investigating the generated language and comparing it to human language Li and Bowling [2019], Lazaridou and Baroni [2020], Lowe et al. [2019]. The research community sees the opportunity to apply this further and scale it up Chaabouni et al. [2022], highlighting the growth possible in this fruitful domain. A key shift this work aims to exploit is moving from message passing being input dependant to task dependant. This raises the level of abstraction of the message from a label to an instruction.

The aim of this research is to utilise the mechanisms of EmCom to make neural network explainability causal. Current xAI methods reveal associative relationships between inputs and outputs but are often limited to analysing single input presentations and say little about causality. The proposed work hinges on the following principles:

• •

  The messages used in EmCom can be made human interpretable (e.g. text or image)

• •

  The message and the output share a causal relationship, as the task can only be solved with the message

• •

  An appropriately generalised language facilitates: transfer learning, data privatisation, distributed learning

All possible truth tables for three inputs ($n=3$) act as the family of tasks that are trained on. This gives $2^{2^{3}}=256$ possible truth tables and $2^{3}*2^{2^{3}}=1024$ training examples. Each truth table is allocated a random One Hot Encoding (OHE) to act as the task ID, this will be passed to the contextualiser to generate a truth table specific message. A random OHE is used to unsure there is as little possible information present in the task ID and all communicated information must be learnt during training. The actor will then receive the message and a set of 3 inputs.

Each agent has a separate contextualiser network and actor network. Both have two hidden layers of 128 ReLU activation neurons, only difference between them being their inputs, outputs, and initialisation. The conextualiser has as many inputs as there are tasks, as it receives a OHE for the task ID. There is a context layer of 32 ReLU neurons places after the final hidden layer. The activation of these neurons forms the message that is passed. There is then a single dense layer which samples from this to predict the expected reward of the actor. The error between this and the actual actor reward forms the reinforcement learning signal. The actor network has $3+32$ inputs (input size + context length) and 2 outputs (binary 0 and 1 for the truth table entry). It is trained using supervised cross-entropy.

Fig. 2 shows how the training data is distributed. Agents are allocated a unique portion of tasks which it can act as a contextualiser for, $r_{c}$. It can act as an actor for all these and an additional portion of the remaining tasks, $r_{a}$. The remaining tasks will be unseen by the agent throughout training. There is a truth table associated with each training example. For each example an agent which can act as a contextualiser is randomly selected. It is given as input the OHE associated with the task ID and produces the associated context. An agent which can be an actor for this training example is then selected and given the associated message and input.

A batch size of 512 is used with an Adam optimiser using learning rate 0.001. Cross-entropy loss between the target and actual output is used to train the actor. Mean-squared error between the predicted and actual actor accuracy is used to train the contextualiser using reinforcement learning. When the contextualiser and actor are both the same agent, gradients are allowed to flow back through the actor and into the communicated message.

The key evaluation metric is the performance of an agent on unseen truth tables. The only way to achieve higher than random chance accuracy is for the communicated messages to contain generalised information about the structure of the family of tasks. This way, even though an agent has never performed the task, it can understand the message and perform the task adequately.

This methodology is currently under development. The principle is to alter the message passing between agents to be human interpretable. This can take the form of an image related to the task, such as the features to look for e.g. the task is looking for a dog so the communicated message contains a nose and floppy ears. This can be injected into the convolutional kernels or concatenated downstream. It is also possible that the message passing can be human language. This would require the message to contain all relevant information to complete the task.

The variable $r_{a}$, controlling the ratio of examples in which an agent is not a contextualiser but can be an actor is explored in Fig, 3. As the amount of ’information sharing’ increases so does the classification accuracy on unseen data. This displays that interaction between agents is key to generating a generalised communication. This is most evident when there is no sharing and the performance does not increase above random chance.