Memory Networks and Transformers for Clinical Reasoning

"The key to clinical reasoning is retrieving the right information at the right time."

August 2026 · Giacomo Saccaggi

---

The Limitation of Fixed Representations

In previous parts, we explored how to represent patients as dense vectors through embeddings and recurrent neural networks. But these approaches share a fundamental limitation: they compress an entire patient history into a fixed-size vector.

Consider a clinician reasoning about a complex case. They don't compress all their knowledge into a single mental state. Instead, they:

• Store relevant facts in working memory
• Query their memory based on the current question
• Retrieve the most relevant pieces of information
• Reason over retrieved facts to form a response

This is exactly what Memory Networks aim to achieve: augmenting neural networks with an external, addressable memory that can be read from and written to.

---

Original Memory Networks (Weston et al., 2015)

The original Memory Network architecture (Weston et al., ICLR 2015) consists of four components that work together to process inputs, store information, and generate responses.

Component I: Input Feature Map

The input component converts raw input $x$ into an internal feature representation $I(x)$. This preprocessing step transforms diverse input types into a common format suitable for memory operations.

For clinical applications, this might encode a patient's chief complaint, a clinical note, or a sequence of diagnosis codes.

Component G: Generalization (Memory Update)

The generalization component updates the memory given a new input. It decides where and how to store new information.

More sophisticated implementations might:

• Forget old memories based on relevance
• Update existing memories rather than creating new ones
• Compress multiple related memories

Component O: Output Feature Map

The output component reads from memory to produce an output representation. This is where the model performs memory addressing — finding the most relevant memories for the current query.

For multi-hop reasoning, the output can be used to query memory again:

Component R: Response

The response component converts the output representation into the final answer format.

The Training Problem: argmax is Not Differentiable

The original Memory Network has a critical limitation: the $\arg\max$ operation in the output component is not differentiable. This means we cannot train the network end-to-end using backpropagation.

This limitation motivated the development of End-to-End Memory Networks.

---

End-to-End Memory Networks (Sukhbaatar et al., 2015)

The breakthrough came with End-to-End Memory Networks, which replace the hard $\arg\max$ with a soft attention mechanism. This makes the entire architecture differentiable and trainable with standard backpropagation.

Architecture Overview

Given a set of inputs $\{x_1, x_2, ..., x_n\}$ to be stored in memory and a query $q$:

1. Memory Encoding: Convert inputs to memory vectors using embedding matrix $A$
2. Query Encoding: Convert query to a vector using embedding matrix $B$
3. Attention: Compute soft attention weights over all memories
4. Output: Weighted sum of memory values (encoded with matrix $C$)
5. Response: Combine with query for final prediction

Three Embedding Matrices

The model uses three separate embedding matrices, each serving a distinct purpose:

This separation of "keys" (what to match on) and "values" (what to retrieve) is a crucial design choice that will reappear in the Transformer architecture.

Soft Attention Mechanism

Instead of selecting a single memory with $\arg\max$, we compute a probability distribution over all memories:

The softmax function is differentiable everywhere, allowing gradients to flow through the memory access operation.

Final Prediction

The output $o$ is combined with the query embedding $u$ to produce the final answer:

The addition of $u$ creates a residual connection, allowing the model to combine retrieved information with the original query.

Multi-hop Reasoning

For complex questions requiring multiple reasoning steps, we can stack multiple "hops" of memory access:

Parameter Sharing Schemes

With multiple hops, we need strategies for sharing parameters across layers:

Scheme	Description	Parameters
Adjacent	Output embedding of layer k becomes input embedding of layer k+1	$A^{k+1} = C^k$
Layer-wise	Same parameters for all layers (like RNN)	$A^1 = A^2 = ... = A^K$
None	Separate parameters per layer	Independent $A^k, C^k$ for each k

Adjacent tying reduces parameters while maintaining layer-specific representations. Layer-wise tying is most parameter-efficient but may limit capacity.

---

Self-Attention: The Foundation of Transformers

Memory Networks use attention to query an external memory. But what if we let each element of a sequence attend to all other elements in the same sequence? This is self-attention, the core mechanism of the Transformer architecture.

The Problem with Sequential Processing

RNNs and LSTMs process sequences one step at a time. This sequential nature creates two problems:

• Limited parallelism: Must wait for previous timesteps before computing current one
• Long-range dependencies: Information must pass through many hidden states, causing gradient issues

Self-attention solves both: every position can attend to every other position in constant path length, and all positions can be computed in parallel.

Query, Key, Value Projections

In self-attention, each input element is projected into three different representations:

This is analogous to End-to-End Memory Networks: Keys correspond to memory embeddings (matrix A), Values to output embeddings (matrix C), and Queries to the query encoding (matrix B).

Scaled Dot-Product Attention

The attention function computes compatibility between queries and keys, then uses these scores to weight values:

The attention matrix $QK^T$ has shape $(n \times n)$, where each entry $(i, j)$ represents how much position $i$ should attend to position $j$.

Multi-Head Attention

Instead of computing attention once, we compute it multiple times in parallel with different learned projections:

Typically, if the model dimension is $d_{model}$ and we use $h$ heads, each head has dimension $d_k = d_v = d_{model}/h$, keeping total computation constant.

---

The Transformer Architecture

The Transformer (Vaswani et al., 2017, "Attention Is All You Need") revolutionized sequence modeling by replacing recurrence entirely with self-attention.

Encoder Block

Each encoder layer consists of two sub-layers:

Key components:

• Residual connections: Add input to output of each sub-layer, enabling gradient flow
• Layer normalization: Stabilizes training by normalizing activations
• Feed-forward network: Applied identically to each position; typically expands then contracts dimensions (e.g., 512 → 2048 → 512)

Decoder Block

The decoder adds a third sub-layer for attending to encoder outputs:

Masked attention prevents positions from attending to subsequent positions during training. This ensures the model can only use information from the past when generating the next token.

Positional Encoding

Self-attention is permutation invariant — it treats inputs as a set, not a sequence. To inject position information, we add positional encodings to the input embeddings:

Alternative: learned positional embeddings, where each position has a trainable embedding vector. Both approaches work similarly in practice.

Why Transformers Excel

Property	RNN/LSTM	Transformer
Maximum path length	$O(n)$	$O(1)$
Computational complexity	$O(n \cdot d^2)$	$O(n^2 \cdot d)$
Parallelization	Sequential	Fully parallel
Long-range dependencies	Difficult (vanishing gradients)	Direct attention

Transformers trade the linear complexity of RNNs for quadratic complexity in sequence length, but gain massive parallelism and better gradient flow.

---

BERT: Pre-trained Bidirectional Transformers

BERT (Bidirectional Encoder Representations from Transformers, Devlin et al., 2018) demonstrated that pre-training Transformers on massive text corpora produces representations that transfer remarkably well to downstream tasks.

From Static to Contextual Embeddings

Previous embedding methods like Word2Vec produce static embeddings: each word has a single vector regardless of context. But words have different meanings in different contexts:

• "The patient has a cold" — illness
• "The room is cold" — temperature
• "He gave me a cold stare" — emotional

BERT produces contextual embeddings: the representation of each word depends on the entire surrounding sentence.

Pre-training Objectives

BERT uses two pre-training tasks:

Fine-tuning for Downstream Tasks

After pre-training, BERT is fine-tuned on task-specific data:

• Classification: Add a classification head on the [CLS] token
• Named Entity Recognition: Add a token-level classifier
• Question Answering: Predict start/end positions in context

---

Healthcare Application: Doctor2Vec

Doctor2Vec (Biswal et al., AAAI 2020) applies memory network principles to learn dynamic doctor representations for clinical trial matching. The key insight: a doctor's expertise should be represented by the patients they have treated, and this representation should be dynamic based on the clinical trial being considered.

The Clinical Trial Matching Problem

Clinical trials need to recruit appropriate physicians to enroll patients. The challenge: how do we match doctors to trials based on their expertise?

• A doctor's expertise is reflected in their patient population
• Different aspects of expertise are relevant for different trials
• Static representations miss trial-specific relevance

Architecture Overview

Doctor2Vec uses a hierarchical memory architecture:

Dynamic Doctor Representation

The key innovation: the doctor embedding changes based on which trial is being considered.

Multimodal Trial Embedding

Clinical trial descriptions contain both structured (eligibility criteria, disease category) and unstructured (text descriptions) information:

Training Objective

The model is trained to predict doctor-trial compatibility:

---

Practical Implementation: End-to-End Memory Network

Let's implement an End-to-End Memory Network for clinical question answering. Given a patient's clinical history (stored as memory) and a clinical question, the model retrieves relevant facts to produce an answer.

Problem Setup

Consider a clinical QA scenario:

# Patient clinical history (facts to store in memory)
facts = [
    "Patient is 65 year old male",
    "Patient has history of type 2 diabetes",
    "Patient takes metformin 500mg twice daily",
    "Latest HbA1c is 7.2%",
    "Patient has hypertension",
    "Patient takes lisinopril 10mg daily",
    "Blood pressure at last visit was 138/85",
    "Patient reports no chest pain",
    "Patient has mild diabetic retinopathy",
]

# Clinical question
question = "What medications is the patient taking?"
# Expected answer: metformin, lisinopril

Memory Network Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class EndToEndMemoryNetwork(nn.Module):
    """End-to-End Memory Network for clinical question answering.
    
    Architecture:
    - Memory: Store clinical facts as key-value pairs
    - Query: Encode the clinical question
    - Attention: Soft attention over memories
    - Multi-hop: Stack multiple attention layers for reasoning
    """
    
    def __init__(
        self,
        vocab_size: int,
        embed_dim: int = 128,
        num_hops: int = 3,
        max_memory_size: int = 50,
        num_classes: int = 100,
    ):
        super().__init__()
        
        self.embed_dim = embed_dim
        self.num_hops = num_hops
        self.max_memory_size = max_memory_size
        
        # Embedding matrices for each hop (using adjacent weight tying)
        # A^k: memory key embeddings, C^k: memory value embeddings
        # With adjacent tying: A^{k+1} = C^k
        self.A = nn.ModuleList()  # Key embeddings
        self.C = nn.ModuleList()  # Value embeddings
        
        for hop in range(num_hops):
            self.A.append(nn.Embedding(vocab_size, embed_dim))
            self.C.append(nn.Embedding(vocab_size, embed_dim))
        
        # Query embedding (B matrix)
        self.B = nn.Embedding(vocab_size, embed_dim)
        
        # Temporal encoding (position embeddings for memory slots)
        self.temporal_A = nn.Parameter(torch.randn(max_memory_size, embed_dim) * 0.1)
        self.temporal_C = nn.Parameter(torch.randn(max_memory_size, embed_dim) * 0.1)
        
        # Output projection
        self.output_layer = nn.Linear(embed_dim, num_classes)
        
        self._init_weights()
    
    def _init_weights(self):
        """Initialize embeddings with small random values."""
        for emb_list in [self.A, self.C]:
            for emb in emb_list:
                nn.init.normal_(emb.weight, mean=0, std=0.1)
        nn.init.normal_(self.B.weight, mean=0, std=0.1)
    
    def forward(
        self,
        memory_tokens: torch.Tensor,   # [batch, num_memories, seq_len]
        memory_mask: torch.Tensor,     # [batch, num_memories] - 1 for real memories
        query_tokens: torch.Tensor,    # [batch, query_len]
    ) -> torch.Tensor:
        """
        Forward pass through the memory network.
        
        Args:
            memory_tokens: Token indices for each memory (clinical fact)
            memory_mask: Binary mask indicating valid memories
            query_tokens: Token indices for the query (clinical question)
            
        Returns:
            logits: [batch, num_classes] classification logits
        """
        batch_size = memory_tokens.size(0)
        num_memories = memory_tokens.size(1)
        
        # Encode query using bag-of-words: sum of word embeddings
        # u: [batch, embed_dim]
        query_embeds = self.B(query_tokens)  # [batch, query_len, embed_dim]
        u = query_embeds.sum(dim=1)  # [batch, embed_dim]
        
        # Multi-hop reasoning
        for hop in range(self.num_hops):
            # Encode memories as keys (what to match on)
            # m: [batch, num_memories, embed_dim]
            mem_key_embeds = self.A[hop](memory_tokens)  # [batch, num_mem, seq_len, embed_dim]
            m = mem_key_embeds.sum(dim=2)  # Bag of words: [batch, num_mem, embed_dim]
            
            # Add temporal encoding
            m = m + self.temporal_A[:num_memories].unsqueeze(0)
            
            # Compute attention: p = softmax(u^T m)
            # [batch, num_memories]
            scores = torch.bmm(m, u.unsqueeze(2)).squeeze(2)
            
            # Mask out padding memories
            scores = scores.masked_fill(memory_mask == 0, float('-inf'))
            p = F.softmax(scores, dim=1)  # [batch, num_memories]
            
            # Encode memories as values (what to retrieve)
            # c: [batch, num_memories, embed_dim]
            mem_val_embeds = self.C[hop](memory_tokens)
            c = mem_val_embeds.sum(dim=2)
            c = c + self.temporal_C[:num_memories].unsqueeze(0)
            
            # Weighted sum of memory values: o = sum_i p_i * c_i
            # [batch, embed_dim]
            o = torch.bmm(p.unsqueeze(1), c).squeeze(1)
            
            # Update query for next hop: u = u + o
            u = u + o
        
        # Final prediction
        logits = self.output_layer(u)  # [batch, num_classes]
        return logits
    
    def get_attention_weights(
        self,
        memory_tokens: torch.Tensor,
        memory_mask: torch.Tensor,
        query_tokens: torch.Tensor,
    ) -> list:
        """Return attention weights for each hop (for interpretability)."""
        batch_size = memory_tokens.size(0)
        num_memories = memory_tokens.size(1)
        
        query_embeds = self.B(query_tokens)
        u = query_embeds.sum(dim=1)
        
        attention_weights = []
        
        for hop in range(self.num_hops):
            mem_key_embeds = self.A[hop](memory_tokens)
            m = mem_key_embeds.sum(dim=2)
            m = m + self.temporal_A[:num_memories].unsqueeze(0)
            
            scores = torch.bmm(m, u.unsqueeze(2)).squeeze(2)
            scores = scores.masked_fill(memory_mask == 0, float('-inf'))
            p = F.softmax(scores, dim=1)
            
            attention_weights.append(p.detach().cpu().numpy())
            
            mem_val_embeds = self.C[hop](memory_tokens)
            c = mem_val_embeds.sum(dim=2)
            c = c + self.temporal_C[:num_memories].unsqueeze(0)
            
            o = torch.bmm(p.unsqueeze(1), c).squeeze(1)
            u = u + o
        
        return attention_weights

Vocabulary and Data Processing

class ClinicalVocabulary:
    """Simple vocabulary for tokenizing clinical text."""
    
    def __init__(self):
        self.word2idx = {'<PAD>': 0, '<UNK>': 1}
        self.idx2word = {0: '<PAD>', 1: '<UNK>'}
        self.next_idx = 2
    
    def add_sentence(self, sentence: str):
        """Add words from a sentence to vocabulary."""
        for word in sentence.lower().split():
            if word not in self.word2idx:
                self.word2idx[word] = self.next_idx
                self.idx2word[self.next_idx] = word
                self.next_idx += 1
    
    def encode(self, sentence: str, max_len: int = 20) -> list:
        """Convert sentence to token indices."""
        tokens = []
        for word in sentence.lower().split()[:max_len]:
            tokens.append(self.word2idx.get(word, 1))  # 1 = <UNK>
        # Pad to max_len
        tokens += [0] * (max_len - len(tokens))
        return tokens
    
    def __len__(self):
        return len(self.word2idx)


def prepare_clinical_data(
    facts: list,
    question: str,
    vocab: ClinicalVocabulary,
    max_memories: int = 20,
    max_seq_len: int = 15,
):
    """Prepare clinical data for the memory network."""
    # Build vocabulary from facts and question
    for fact in facts:
        vocab.add_sentence(fact)
    vocab.add_sentence(question)
    
    # Encode memories (clinical facts)
    memory_tokens = []
    for fact in facts[:max_memories]:
        memory_tokens.append(vocab.encode(fact, max_seq_len))
    
    # Pad to max_memories
    num_real_memories = len(memory_tokens)
    while len(memory_tokens) < max_memories:
        memory_tokens.append([0] * max_seq_len)
    
    # Create memory mask
    memory_mask = [1] * num_real_memories + [0] * (max_memories - num_real_memories)
    
    # Encode query
    query_tokens = vocab.encode(question, max_seq_len)
    
    return {
        'memory_tokens': torch.tensor([memory_tokens]),
        'memory_mask': torch.tensor([memory_mask]),
        'query_tokens': torch.tensor([query_tokens]),
    }

Example Usage

# Clinical case
facts = [
    "Patient is 65 year old male",
    "Patient has history of type 2 diabetes",
    "Patient takes metformin 500mg twice daily",
    "Latest HbA1c is 7.2 percent",
    "Patient has hypertension",
    "Patient takes lisinopril 10mg daily",
    "Blood pressure at last visit was 138 over 85",
    "Patient reports no chest pain",
    "Patient has mild diabetic retinopathy",
]

question = "What medications is the patient taking"

# Prepare data
vocab = ClinicalVocabulary()
data = prepare_clinical_data(facts, question, vocab)

# Initialize model
model = EndToEndMemoryNetwork(
    vocab_size=len(vocab),
    embed_dim=64,
    num_hops=3,
    max_memory_size=20,
    num_classes=10,  # Number of possible answers
)

# Forward pass
with torch.no_grad():
    logits = model(
        data['memory_tokens'],
        data['memory_mask'],
        data['query_tokens'],
    )
    
    # Get attention weights for interpretability
    attention = model.get_attention_weights(
        data['memory_tokens'],
        data['memory_mask'],
        data['query_tokens'],
    )

# Visualize attention over clinical facts
print("Question:", question)
print("\nAttention weights per hop:")
for hop, weights in enumerate(attention):
    print(f"\nHop {hop + 1}:")
    for i, (fact, w) in enumerate(zip(facts, weights[0])):
        if w > 0.05:  # Only show significant attention
            print(f"  [{w:.3f}] {fact}")

# Sample output (after training):
# Question: What medications is the patient taking
# 
# Attention weights per hop:
# 
# Hop 1:
#   [0.342] Patient takes metformin 500mg twice daily
#   [0.298] Patient takes lisinopril 10mg daily
#   [0.112] Patient has history of type 2 diabetes
# 
# Hop 2:
#   [0.456] Patient takes metformin 500mg twice daily
#   [0.389] Patient takes lisinopril 10mg daily
# 
# Hop 3:
#   [0.521] Patient takes metformin 500mg twice daily
#   [0.423] Patient takes lisinopril 10mg daily

---

Practical Implementation: Transformer for EHR Classification

Now let's implement a Transformer encoder for classifying patient EHR sequences. Given a sequence of visits (each represented by diagnosis codes), we predict disease risk.

Positional Encoding

import math

class PositionalEncoding(nn.Module):
    """Sinusoidal positional encoding for sequences.
    
    Adds position-dependent signals so the model can learn to attend
    by relative position, not just content similarity.
    """
    
    def __init__(self, d_model: int, max_len: int = 500, dropout: float = 0.1):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        
        # Create positional encoding matrix
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        
        # Compute the div_term for sinusoidal encoding
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )
        
        # Even indices: sin, Odd indices: cos
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        
        # Register as buffer (not a parameter, but saved with model)
        pe = pe.unsqueeze(0)  # [1, max_len, d_model]
        self.register_buffer('pe', pe)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: [batch, seq_len, d_model]
        Returns:
            x + positional encoding
        """
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)

Multi-Head Self-Attention

class MultiHeadAttention(nn.Module):
    """Multi-head self-attention mechanism.
    
    Computes attention over input sequence, allowing each position
    to attend to all other positions with learned query/key/value projections.
    """
    
    def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1):
        super().__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        # Linear projections for Q, K, V
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        
        # Output projection
        self.W_o = nn.Linear(d_model, d_model)
        
        self.dropout = nn.Dropout(p=dropout)
        self.scale = math.sqrt(self.d_k)
    
    def forward(
        self,
        x: torch.Tensor,
        mask: torch.Tensor = None,
    ) -> torch.Tensor:
        """
        Args:
            x: [batch, seq_len, d_model]
            mask: [batch, seq_len] - 1 for valid positions, 0 for padding
        Returns:
            output: [batch, seq_len, d_model]
        """
        batch_size, seq_len, _ = x.size()
        
        # Project to Q, K, V
        Q = self.W_q(x)  # [batch, seq_len, d_model]
        K = self.W_k(x)
        V = self.W_v(x)
        
        # Reshape for multi-head attention: [batch, num_heads, seq_len, d_k]
        Q = Q.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        
        # Scaled dot-product attention
        # scores: [batch, num_heads, seq_len, seq_len]
        scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale
        
        # Apply mask if provided (for padding)
        if mask is not None:
            # Expand mask: [batch, 1, 1, seq_len]
            mask = mask.unsqueeze(1).unsqueeze(2)
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        # Softmax over keys
        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)
        
        # Apply attention to values
        # output: [batch, num_heads, seq_len, d_k]
        output = torch.matmul(attn_weights, V)
        
        # Reshape back: [batch, seq_len, d_model]
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
        
        # Final projection
        output = self.W_o(output)
        
        return output

Transformer Encoder Block

class TransformerEncoderBlock(nn.Module):
    """Single Transformer encoder block.
    
    Architecture:
    - Multi-head self-attention
    - Add & Normalize (residual connection + layer norm)
    - Feed-forward network
    - Add & Normalize
    """
    
    def __init__(
        self,
        d_model: int,
        num_heads: int,
        d_ff: int,
        dropout: float = 0.1,
    ):
        super().__init__()
        
        # Multi-head attention
        self.attention = MultiHeadAttention(d_model, num_heads, dropout)
        
        # Feed-forward network
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
        )
        
        # Layer normalization
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        
        # Dropout for residual connections
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
        """
        Args:
            x: [batch, seq_len, d_model]
            mask: [batch, seq_len] padding mask
        Returns:
            output: [batch, seq_len, d_model]
        """
        # Self-attention with residual connection
        attn_output = self.attention(x, mask)
        x = self.norm1(x + self.dropout(attn_output))
        
        # Feed-forward with residual connection
        ffn_output = self.ffn(x)
        x = self.norm2(x + self.dropout(ffn_output))
        
        return x

Complete Transformer for EHR Classification

class EHRTransformer(nn.Module):
    """Transformer encoder for EHR sequence classification.
    
    Takes a sequence of patient visits, each represented by diagnosis code
    embeddings, and predicts disease risk or other outcomes.
    """
    
    def __init__(
        self,
        vocab_size: int,
        d_model: int = 128,
        num_heads: int = 4,
        num_layers: int = 2,
        d_ff: int = 256,
        max_seq_len: int = 100,
        num_classes: int = 2,
        dropout: float = 0.1,
    ):
        super().__init__()
        
        # Code embedding: convert diagnosis codes to vectors
        self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=0)
        
        # Positional encoding for visit ordering
        self.pos_encoding = PositionalEncoding(d_model, max_seq_len, dropout)
        
        # Stack of Transformer encoder blocks
        self.encoder_layers = nn.ModuleList([
            TransformerEncoderBlock(d_model, num_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])
        
        # Classification head
        self.classifier = nn.Sequential(
            nn.Linear(d_model, d_model // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_model // 2, num_classes),
        )
        
        self.d_model = d_model
    
    def forward(
        self,
        visit_codes: torch.Tensor,     # [batch, num_visits, codes_per_visit]
        visit_mask: torch.Tensor,       # [batch, num_visits] - 1 for real visits
    ) -> torch.Tensor:
        """
        Args:
            visit_codes: Diagnosis codes for each visit
            visit_mask: Mask for padding visits
        Returns:
            logits: [batch, num_classes]
        """
        batch_size, num_visits, codes_per_visit = visit_codes.size()
        
        # Embed codes: [batch, num_visits, codes_per_visit, d_model]
        code_embeds = self.embedding(visit_codes)
        
        # Aggregate codes within each visit (sum pooling)
        # [batch, num_visits, d_model]
        visit_embeds = code_embeds.sum(dim=2)
        
        # Scale by sqrt(d_model) for stability
        visit_embeds = visit_embeds * math.sqrt(self.d_model)
        
        # Add positional encoding
        visit_embeds = self.pos_encoding(visit_embeds)
        
        # Apply Transformer encoder layers
        x = visit_embeds
        for encoder in self.encoder_layers:
            x = encoder(x, visit_mask)
        
        # Global pooling: average over non-padded visits
        # [batch, d_model]
        mask_expanded = visit_mask.unsqueeze(-1).float()
        x_masked = x * mask_expanded
        pooled = x_masked.sum(dim=1) / mask_expanded.sum(dim=1).clamp(min=1)
        
        # Classification
        logits = self.classifier(pooled)
        
        return logits

Training Example

# Example: Predict heart disease risk from visit sequence

# Sample data: visits with diagnosis codes (indices)
# Visit 1: [Diabetes, Hypertension]
# Visit 2: [Hyperlipidemia, Obesity]
# Visit 3: [Chest pain, ECG abnormal]

# Create sample batch
batch_size = 4
num_visits = 10
codes_per_visit = 5
vocab_size = 1000

# Random data for demonstration
visit_codes = torch.randint(1, vocab_size, (batch_size, num_visits, codes_per_visit))
visit_mask = torch.ones(batch_size, num_visits)
visit_mask[:, 7:] = 0  # Last 3 visits are padding

labels = torch.randint(0, 2, (batch_size,))  # Binary: heart disease yes/no

# Initialize model
model = EHRTransformer(
    vocab_size=vocab_size,
    d_model=128,
    num_heads=4,
    num_layers=2,
    d_ff=256,
    max_seq_len=50,
    num_classes=2,
    dropout=0.1,
)

# Training setup
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Single training step
model.train()
optimizer.zero_grad()

logits = model(visit_codes, visit_mask)
loss = criterion(logits, labels)
loss.backward()
optimizer.step()

print(f"Loss: {loss.item():.4f}")
print(f"Predictions: {torch.argmax(logits, dim=1).tolist()}")
print(f"Labels: {labels.tolist()}")

# Output:
# Loss: 0.7123
# Predictions: [1, 0, 1, 0]
# Labels: [0, 1, 1, 0]

Model Size Comparison

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

print(f"EHR Transformer parameters: {count_parameters(model):,}")

# Compare with Memory Network
memnet = EndToEndMemoryNetwork(
    vocab_size=vocab_size,
    embed_dim=128,
    num_hops=3,
    num_classes=2,
)
print(f"Memory Network parameters: {count_parameters(memnet):,}")

# Output:
# EHR Transformer parameters: 424,962
# Memory Network parameters: 898,306

---

Comparing Architectures for Clinical Reasoning

Each architecture we've covered has different strengths for healthcare applications:

Architecture	Memory	Attention	Best For
Original Memory Networks	Explicit external	Hard (argmax)	Strongly supervised QA
End-to-End MemNets	Explicit external	Soft (softmax)	Clinical fact retrieval, multi-hop reasoning
Transformer Encoder	Implicit (in weights)	Self-attention	Sequential EHR modeling, parallel processing
BERT	Pre-trained knowledge	Bidirectional self-attention	Clinical NLP, note understanding
Doctor2Vec	Patient history as memory	Dynamic query-based	Doctor-trial matching, personalized retrieval

When to Use What

Task	Recommended Architecture	Reasoning
Clinical Question Answering	Memory Network	Explicit retrieval from patient facts
Disease Prediction from Visits	Transformer	Captures temporal patterns in visit sequences
Clinical Note Classification	Clinical BERT	Leverages pre-trained medical language understanding
Named Entity Recognition	BERT + Token Classifier	Context-aware token representations
Doctor Recommendation	Memory Network (Doctor2Vec)	Dynamic expertise representation
Multi-modal EHR + Notes	Transformer Fusion	Handles heterogeneous inputs with attention

---

Key Takeaways

---

Summary

Concept	Key Idea	Clinical Application
Memory Networks	External memory + attention for retrieval	Clinical QA, fact retrieval
End-to-End Training	Soft attention replaces hard argmax	Learning without explicit supervision
Self-Attention	Every position attends to all others	Visit sequence modeling
Transformer	Stack of self-attention + FFN layers	EHR classification, parallel processing
Positional Encoding	Inject position information	Capture visit ordering in patient history
BERT	Pre-trained bidirectional Transformer	Clinical note understanding
Doctor2Vec	Patient memory + trial query attention	Clinical trial matching

These architectures represent the evolution from external memory augmentation to fully self-contained attention mechanisms. Memory Networks excel when we have structured facts to retrieve; Transformers excel when we need to model complex relationships within sequences; and BERT excels when we have abundant unlabeled text for pre-training. In healthcare, the choice depends on your data modality (structured codes vs. clinical text), task requirements (retrieval vs. classification), and interpretability needs.

---

References

• Weston, Chopra & Bordes (2015). "Memory Networks" — Original Memory Networks, ICLR 2015
• Sukhbaatar et al. (2015). "End-To-End Memory Networks" — Differentiable memory access, NeurIPS 2015
• Vaswani et al. (2017). "Attention Is All You Need" — Transformer architecture, NeurIPS 2017
• Devlin et al. (2018). "BERT: Pre-training of Deep Bidirectional Transformers" — BERT, NAACL 2019
• Biswal et al. (2020). "Doctor2Vec: Dynamic Doctor Representation Learning" — AAAI 2020
• Alsentzer et al. (2019). "Publicly Available Clinical BERT Embeddings" — ClinicalBERT
• Lee et al. (2020). "BioBERT: a pre-trained biomedical language representation model"