AI for Healthcare: Medical Imaging and Drug Discovery

Two data structures, one underlying question

A radiologist at a tertiary cancer centre reading a glioma MRI is not asking a yes/no question. A classifier that outputs "tumour present, 94% confidence" is close to useless to the radiation oncologist who has to plan the beam. Radiotherapy for a brain tumour delivers a lethal dose to a tightly bounded 3-D region and must spare healthy tissue a few millimetres away — the optic nerve, the brainstem. The clinical question is not "is there a tumour" but "which exact voxels are tumour," because that boundary becomes the physical edge of where a radiation beam is and is not allowed to go. This is why medical imaging AI is dominated by segmentation architectures, not the plain image classifiers you built in Grade 11 — and why the field's workhorse architecture, U-Net, looks structurally different from a ResNet or VGG-style classifier even though both are convolutional networks.

A few kilometres away in a computational chemistry lab, a completely different data problem shows up. A drug candidate is a molecule — atoms connected by bonds — and a model asked to predict "will this molecule bind this protein pocket" or "is this molecule toxic" cannot use a convolutional kernel at all, for a structural reason: convolution assumes a fixed grid where every position has the same number of neighbours in the same geometric arrangement (a pixel always has up to 8 neighbours in a 3×3 kernel). A molecule has no grid. A benzene ring has six neighbours arranged in a cycle; a single terminal —OH has one. There is no canonical way to list a molecule's atoms in a fixed-length array, because the same molecule written from two different starting atoms produces two differently-ordered atom lists that must nonetheless yield the identical prediction. Solving this requires an architecture built for graphs, not grids: the message-passing graph neural network (GNN).

Both problems reduce to the same underlying question that recurs across deep learning: what is the correct inductive bias — the built-in structural assumption — for this data's actual geometry? Get the inductive bias right (grid-preserving encoder-decoder for pixels, permutation-invariant message passing for graphs) and a comparatively small, interpretable model matches specialist performance. Get it wrong and no amount of scale fixes it. This chapter works through both mechanisms end to end, with fully traced numeric examples, and then shows where they meet in a real production system: DeepMind's AlphaFold2.

Mechanism 1 — U-Net: why segmentation needs an encoder AND a decoder

A standard classification CNN (the kind you built for MNIST or CIFAR-style tasks) repeatedly convolves and pools an image, shrinking its spatial resolution while growing its channel depth, until the spatial dimensions collapse to a single vector fed into a softmax. That downsampling path is exactly what you want for "what is in this image" — pooling builds translation-invariant, high-level semantic features layer by layer. But it is exactly wrong for "where is it," because every pooling operation throws away spatial precision. By the time a classifier reaches its bottleneck, it has "seen" a tumour-shaped blob somewhere in the image but has lost the pixel-exact boundary of that blob permanently — pooling is not invertible.

Ronneberger, Fischer, and Brox's 2015 MICCAI paper "U-Net: Convolutional Networks for Biomedical Image Segmentation" solves this with a symmetric encoder-decoder: the same downsampling path as a classifier (the "encoder," answering what), followed by a mirrored upsampling path (the "decoder," reconstructing spatial resolution back to the input size) — plus the architecture's defining trick, skip connections that carry each encoder stage's full-resolution feature map directly across to the matching decoder stage, concatenated channel-wise. The decoder alone, working only from the compressed bottleneck, could reconstruct roughly where a lesion is but not its exact edge — that fine-grained boundary information was destroyed by pooling and cannot be regenerated from a compressed representation. The skip connection re-injects it: at each decoder stage, the network combines the deep, semantic "this region is tumour tissue" signal coming up from the bottleneck with the shallow, precise "this exact pixel's local texture" signal arriving directly from the corresponding encoder layer. Segmentation quality comes from that fusion, not from either path alone.

There is a second, equally deliberate design choice: the loss function. A brain tumour frequently occupies under 5% of an MRI slice's pixels. A model trained with plain per-pixel cross-entropy can achieve deceptively low loss by predicting "background" almost everywhere, because 95% of pixels genuinely are background — the loss is dominated by the class that is easiest to get right and barely penalized for missing the lesion entirely. Segmentation networks are instead typically trained with a Dice loss (1 − Dice coefficient), which measures overlap directly and is insensitive to this class imbalance because it never counts true negatives at all.

Worked example: computing Dice and IoU by hand

Take a toy 4×4 slice. The ground-truth tumour mask (1 = tumour, 0 = background) and a trained model's raw sigmoid probability output are:

import numpy as np

ground_truth = np.array([
    [0, 0, 1, 1],
    [0, 1, 1, 1],
    [0, 1, 1, 0],
    [0, 0, 0, 0],
])

prediction_probs = np.array([
    [0.10, 0.20, 0.85, 0.60],
    [0.15, 0.40, 0.90, 0.55],
    [0.55, 0.65, 0.70, 0.20],
    [0.05, 0.10, 0.15, 0.05],
])

prediction = (prediction_probs >= 0.5).astype(int)

intersection = np.sum(ground_truth * prediction)
gt_sum = np.sum(ground_truth)
pred_sum = np.sum(prediction)

dice = 2 * intersection / (gt_sum + pred_sum)
iou = intersection / (gt_sum + pred_sum - intersection)

print(prediction)
print(intersection, gt_sum, pred_sum)
print(dice, iou)

Trace it by hand, row by row of prediction_probs, thresholding at 0.5: row 0 gives [0, 0, 1, 1]; row 1 gives [0, 0, 1, 1] — note position (1,1) is 0.40, below threshold, even though the ground truth has a 1 there, a false negative; row 2 gives [1, 1, 1, 0] — position (2,0) is 0.55, above threshold, even though the ground truth has a 0 there, a false positive; row 3 gives [0, 0, 0, 0]. So prediction = [[0,0,1,1],[0,0,1,1],[1,1,1,0],[0,0,0,0]].

Now count. ground_truth has 7 ones (2 + 3 + 2 + 0). prediction also has 7 ones (2 + 2 + 3 + 0) — coincidentally equal totals despite one false positive and one false negative cancelling out numerically. The intersection (both arrays 1 at the same cell) is 6: two hits in row 0, two in row 1 (missing the false-negative cell), two in row 2 (the false-positive cell at (2,0) doesn't count, since ground truth is 0 there), zero in row 3.

Dice = 2 × 6 / (7 + 7) = 12/14 = 6/7 ≈ 0.857. IoU = 6 / (7 + 7 − 6) = 6/8 = 0.75. The code prints exactly these two numbers. Notice Dice > IoU here — this is not a coincidence of this example but an algebraic identity: Dice = 2·IoU / (1 + IoU) always, for any segmentation. You can verify it right now: 2 × 0.75 / 1.75 = 1.5/1.75 = 0.857142..., matching. This identity is a useful sanity check on any segmentation metrics table you're handed — if reported Dice and IoU don't satisfy it, one of the two numbers was computed wrong.

Common misconception: "high accuracy means the model learned the pathology"

A student's natural assumption is that if a CNN reaches 95%+ accuracy or a high Dice score on a disease-detection task, it must have learned to recognize the actual radiological signs of that disease the way a trained radiologist does. This is false in a specific, documented way, and the failure mode has a name: shortcut learning. A convolutional network's loss function only rewards whatever correlates with the label in the training distribution — it has no preference for causally meaningful features over spuriously correlated ones, and spurious features are frequently easier to fit.

Zech, Badgeley, Liu, Costa, Titano, and Oermann's 2018 PLOS Medicine study, "Variable generalization performance of a deep learning model to detect pneumonia in chest radiographs," trained CNNs to detect pneumonia on chest X-rays from multiple hospital systems. A model trained and tested within the same hospital scored a high AUC. But performance dropped sharply when the identical model was tested on X-rays from a different hospital it hadn't trained on — the classic signature of a model that memorized dataset-specific artifacts rather than the disease. The mechanism the paper identified: portable chest X-rays (taken at a patient's bedside, typically for sicker patients) carry visible equipment and positioning differences — burned-in markers, laterality tokens, image characteristics tied to which specific hospital's specific portable X-ray machine took the scan — and pneumonia prevalence itself varies by hospital and by which patients get portable versus standard films. The network could get a high training-set AUC by learning "this looks like a portable film from Hospital A, where pneumonia is more common" rather than by learning the actual radiological pattern of consolidated lung tissue. That shortcut works perfectly on data from the training hospital and collapses on any other hospital's equipment.

The corrected mental model: a segmentation or classification network's accuracy number tells you only how well it fits correlations present in its training distribution — nothing about which features it used to do so. Before trusting any medical imaging model's reported accuracy, you check generalization to genuinely external data (different hospital, different scanner, different population) and inspect what the model is actually attending to, because a number alone cannot distinguish "learned the pathology" from "learned a shortcut that happens to correlate with the pathology in this dataset."

Mechanism 2 — message passing: learning on molecules instead of pixels

A drug candidate's molecular graph has nodes (atoms, each with a type — carbon, oxygen, nitrogen) and edges (bonds, each with a type — single, double, aromatic). Two requirements a working architecture must satisfy that a grid convolution does not: it must handle a variable number of neighbours per node (an atom's "degree" — how many bonds it has — ranges from 1 to 4 or more, unlike a pixel's fixed 8 grid-neighbours), and it must be permutation-invariant — the exact same molecule, with its atoms listed in a different order in the input file, must produce the identical prediction, because atom ordering carries no chemical meaning.

The message-passing neural network framework — formalized by Gilmer, Schoenholz, Riley, Vinyals, and Dahl in their 2017 ICML paper "Neural Message Passing for Quantum Chemistry," building on Duvenaud et al.'s 2015 NeurIPS paper "Convolutional Networks on Graphs for Learning Molecular Fingerprints" — solves both at once. Each atom starts with a feature vector (encoding its element type, charge, and so on). In one "round" of message passing, every node computes an update using only (a) its own current feature vector and (b) the sum (an order-independent operation, unlike concatenation) of its direct bonded neighbours' current feature vectors. Because the aggregation step is a sum over whatever neighbours actually exist, the same update rule works identically whether a node has one bond or four, and reordering the neighbour list changes nothing, since addition is commutative.

Worked example: one message-passing round, traced by hand

Take a 3-atom fragment forming a path graph, A — B — C, with initial scalar embeddings h_A = 1.0, h_B = 1.0, h_C = 3.0 (imagine A and B as carbons and C as a more electronegative atom like oxygen — the exact numbers are illustrative, chosen so the arithmetic traces cleanly). The update rule, a simplified one-layer version of the Gilmer et al. framework: h′_v = ReLU(0.5·h_v + 0.5·Σ_{u ∈ N(v)} h_u), where N(v) is v's set of bonded neighbours.

def relu(x):
    return max(x, 0)

h = {'A': 1.0, 'B': 1.0, 'C': 3.0}
edges = {'A': ['B'], 'B': ['A', 'C'], 'C': ['B']}
W_self, W_msg = 0.5, 0.5

h_new = {}
for v in h:
    m_v = sum(h[u] for u in edges[v])
    h_new[v] = relu(W_self * h[v] + W_msg * m_v)

print(h_new)

Trace each node. A's only neighbour is B, so its message is m_A = h_B = 1.0, giving h′_A = ReLU(0.5×1.0 + 0.5×1.0) = ReLU(1.0) = 1.0. B has two neighbours, A and C, so m_B = h_A + h_C = 1.0 + 3.0 = 4.0, giving h′_B = ReLU(0.5×1.0 + 0.5×4.0) = ReLU(0.5 + 2.0) = 2.5. C's only neighbour is B, so m_C = h_B = 1.0, giving h′_C = ReLU(0.5×3.0 + 0.5×1.0) = ReLU(1.5 + 0.5) = 2.0. The code prints {'A': 1.0, 'B': 2.5, 'C': 2.0}, matching exactly.

The pedagogically important result is B's value: it moved further from its starting value (1.0 → 2.5) than either terminal atom did, purely because of its position in the graph — it has two neighbours pulling its embedding, not one. This is the entire point of a GNN over a naive "bag of atoms" baseline that would encode each atom independently of its bonding context: the same carbon atom gets a different learned embedding depending on what it's connected to, which is exactly the structural information that determines how a molecule actually behaves — its shape, polarity, and where it can bind a protein pocket. A final readout step pools all updated node embeddings into one molecule-level vector (commonly by summing, for the same permutation-invariance reason): here, 1.0 + 2.5 + 2.0 = 5.5, which a downstream linear layer would map to a predicted property — binding affinity, solubility, toxicity. In a real virtual-screening pipeline, this scoring step runs across millions of candidate molecules to rank which few thousand are worth synthesizing and testing in a wet lab, long before any question of clinical trials arises.

Where the two mechanisms meet: AlphaFold2

Structure-based drug design needs a 3-D model of the target protein's binding pocket before any ligand can be screened against it — and for most of biology's history, that structure came only from slow, expensive experimental methods (X-ray crystallography, cryo-EM). Jumper, Evans, Pritzel, and colleagues at DeepMind published "Highly accurate protein structure prediction with AlphaFold" in Nature in 2021, reporting prediction accuracy competitive with experimental methods across the CASP14 benchmark. AlphaFold2's core module, the Evoformer, operates on two representations simultaneously: a multiple sequence alignment (MSA) representation capturing evolutionary co-variation across related protein sequences, and a pairwise representation capturing the network's evolving belief about the spatial relationship between every pair of residues. These two representations are repeatedly updated by attention — the same self-attention mechanism from the transformer architecture you studied for language, here applied row-wise and column-wise across the MSA and with triangular consistency updates across the pair representation, so information about residue i and j's relationship can also constrain the network's belief about j and k. A final structure module then converts this refined pairwise representation into actual 3-D atomic coordinates.

The connection to both mechanisms above: attention over a set of residues with pairwise interactions is itself a form of message passing on a graph (here, a graph where every residue can in principle message every other residue, not just covalently bonded neighbours) — and the predicted protein structure is the input a GNN-based ligand-screening pipeline needs before it can meaningfully score candidate molecules' binding affinity against that specific pocket. But a predicted structure is a single static snapshot; a real protein flexes, and a real binding event is a dynamic process. AlphaFold2's output functions as a strong starting hypothesis that feeds into further modelling — docking simulations, molecular dynamics — not as a final, standalone answer about how a drug will actually behave in a binding pocket.

Active recall

Attempt each question before reading its answer.

1. In the Dice/IoU worked example, suppose the false positive at pixel (2,0) is corrected — the model no longer predicts tumour there — while everything else about the prediction stays exactly the same. Recompute Dice and IoU, and state whether precision, recall, both, or neither improved.

2. Why can't a standard 2-D convolutional kernel, as used in a U-Net's encoder, be applied directly to a molecular graph the way it's applied to an image?

3. Extend the A–B–C fragment by attaching a fourth atom D to C (chain A–B–C–D), with h_D = 1.0. Recompute h′_C after one message-passing round with this new edge. Does h′_A or h′_B change in that same round? Explain using the concept of receptive field.

4. True or false, with justification: a segmentation model reporting Dice = 0.857 on some slice must also have IoU = 0.857 on that same slice.

5. AlphaFold2 predicts a single static 3-D structure, but a real protein's binding behaviour is dynamic. What does this imply about how a structure prediction should be used in early-stage, structure-based drug discovery?

6. A pneumonia-detection model trained on portable chest X-rays from one hospital's ICU reaches high AUC in-house but degrades sharply when tested at a different hospital. Using the shortcut-learning mechanism from this chapter (not a demographic-representation argument), give one plausible technical cause, and one concrete way to test for it before deployment.

Answers

1. The new prediction row 2 becomes [0, 1, 1, 0] instead of [1, 1, 1, 0], so pred_sum drops from 7 to 6. Intersection is unchanged at 6, because (2,0) was never counted in the intersection in the first place — ground truth was 0 there. New Dice = 2×6/(7+6) = 12/13 ≈ 0.923. New IoU = 6/(7+6−6) = 6/7 ≈ 0.857. Precision = intersection/pred_sum: it rises from 6/7 ≈ 0.857 to 6/6 = 1.0, a real improvement, since removing a false positive directly raises precision. Recall = intersection/gt_sum stays at 6/7 ≈ 0.857 unchanged, since neither the intersection nor the ground-truth count moved — the false negative at (1,1) is untouched by this edit. So precision improved, recall did not.

2. A 2-D convolution assumes a fixed, regular grid where every position has a consistent, fixed-size neighbourhood (e.g. 8 neighbours in a 3×3 kernel) and a canonical spatial ordering (up/down/left/right are always the same directions). A molecular graph has neither property: atoms have variable degree (1 to 4+ bonds), and there is no canonical ordering of atoms — the same molecule listed in a different atom order in an input file is chemically identical and must produce an identical prediction. A grid convolution applied naively would treat two orderings of the same molecule as different inputs and would have no consistent way to handle atoms with different numbers of neighbours. Message-passing GNNs solve this by aggregating over each node's actual (variable-size) neighbour set with an order-invariant operation like summation.

3. New edges: C is now bonded to both B and D. h′_C = ReLU(0.5×h_C + 0.5×(h_B + h_D)) = ReLU(0.5×3.0 + 0.5×(1.0 + 1.0)) = ReLU(1.5 + 1.0) = 2.5 — changed from the original 2.0, because C's neighbourhood changed. h′_A and h′_B are unaffected in this same round: A's formula only ever references h_A and h_B (unchanged values), and B's formula only references h_A and h_C's pre-update values (also unchanged — message passing at round 1 always uses round-0 values for every node, before any node has been updated). This is the graph analogue of a CNN's receptive field: one message-passing round only lets information travel exactly one hop, so a structural change at D is invisible to nodes more than one hop away until enough additional rounds have run — a second round begins to move B (2 hops from D), and only a third round begins to move A (3 hops from D). Stacking more message-passing layers, exactly like stacking more convolutional layers, grows how far information can propagate through the structure.

4. False. Dice and IoU are related by the identity Dice = 2·IoU/(1+IoU), which rearranges to IoU = Dice/(2 − Dice); they are equal only in the degenerate case of perfect overlap (both = 1). Plugging in Dice = 6/7 ≈ 0.857: IoU = 0.857/(2 − 0.857) = 0.857/1.143 ≈ 0.75, matching the worked example exactly. Dice is always ≥ IoU for the same prediction (except at 1.0), so a claim of equal Dice and IoU below 1.0 signals a computation error.

5. A predicted structure should be treated as a starting hypothesis feeding further computational steps — docking simulations against candidate ligands, molecular dynamics to explore conformational flexibility, experimental validation — not as a finished, standalone answer about how the protein will behave when a drug binds it. Binding is a dynamic process involving conformational change; a single static coordinate set cannot capture that by itself.

6. A plausible cause: shortcut learning on hospital-specific, non-pathological artifacts — burned-in equipment markers, laterality tokens, or other portable-X-ray-machine-specific image characteristics that happened to correlate with pneumonia prevalence at the training hospital, exactly the failure mode Zech et al. documented, rather than the model having learned the actual radiological signs of consolidation. A concrete pre-deployment test: evaluate the model on a genuinely external hold-out set from a hospital and scanner it never trained on (not just a held-out split of the same hospital's data), and separately generate saliency or Grad-CAM maps to check whether the model's attention concentrates on anatomically plausible lung regions rather than image borders, corners, or burned-in text.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where ai for healthcare: medical imaging and drug discovery is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting ai for healthcare: medical imaging and drug discovery to at least 3 other topics you have studied.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind ai for healthcare: medical imaging and drug discovery, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.