Approaches to biomedical knowledge

Overview

• BERT: training
• fine tuning
• sentence classification
• NER
• SBERT

BERT

• BERT was trained on Wikipedia (~2.5B words) and Google’s BooksCorpus (~800M words).
• Trained at Google on 64 TPUs (Tensor Processing Units - roughly speaking, Google’s version of the GPU) over 4 days
• Many other versions of BERT have been trained for various purposes

Causal vs masked language modeling

• Predecessors modelled language unidirectionally – train on predicting the next word
• BERT is designed to extract bidirectional representations: predict the masked word given both left and right context

BERT: unidirectional vs bidirectional

• The authors argue that the unidirectionality of models such as GPT limit the choice of architectures that can be used during pre-training because every token can only attend to previous tokens in the self-attention layers of the Transformer.
• Such restrictions are sub-optimal for sentence-level tasks, and could be very harmful when applying finetuning based approaches to token-level tasks such as question answering, where it is crucial to incorporate context from both directions

BERT training

• BERT is an encoder-only transformer and the training is similar to what was described in previous lectures.

BERT Architecture

BERT vs. GPT

• We will provide some examples below

Masked Language Modeling (MLM)

• BERT is forced to utilize words on either side of the token simultaneously to predict the missing mask.
• Comparing bidirectional reasoning capabilities across different sentence structures:

Left and right context in BERT

Recall the definition of the attention output

\[ \underbrace{\mathrm{Attention}(\mathbf{Q},\mathbf{K},\mathbf{V})}_{N\times N} = \mathrm{softmax}\left( \frac{ \underbrace{\mathbf{Q}}_{N\times d_{hidden}} \times \underbrace{\mathbf{K}^T}_{d_{hidden}\times N}} {\sqrt{d_{k}}} \right) \]

• Recall that for the masked attention in the decoder, we mask words to “zero out” interactions of tokens with other subsequent tokens

Left and right context in BERT

• In contrast, with BERT there is no mask
• Each token therefore attends to other tokens to the left and the right
• The pretraining procedure instead selects 15% of the tokens from the sentence to be masked
  • 80% of the time, the token is replaced by [MASK]: Berlin is the [MASK] of Germany
  • 10% of the time, the token is replaced by a random word: Berlin is the hedgehog of Germany
  • 10% of the time, the token is not replaced: Berlin is the capital of Germany

Loss: BERT

• Like any encoder, the BERT model produces embeddings as output.
• We focus on the output for the masked position to calculate the loss
• Backpropation and gradient descent are as usual

Loss: BERT

• In detail,
• The LM head takes the output of the final transformation layer \(L\), multiplies it by the unemedding layer and uses softmax to transform the output into probabilities

\[ \begin{align} \mathbf{u}_i &= \mathbf{h}_i^L \mathbf{E}^T \\ \mathbf{y}_i &= \mathrm{softmax}(\mathbf{u}_i) \end{align} \]

• In our example, the \(x_i\) that corresponds to “capital”, the loss is the probability of the correct word “capital” given the output \[ L_{MLM}(x_i) = -\log P(x_i\mid \mathbf{h}_i^L) \]
• Gradients are obtained by taking the average loss over a batch

\[ -\frac{1}{M}\sum_{m=1}^M \log P(x_i\mid \mathbf{h}_i^L) \]

Next-sentence prediction (NSP)

• BERT is therefore also trained on the next-sentence prediction task
• 50% of the time, correct pairs of sentence are chosen

Why did the chicken cross the road? To get to the other side

• 50% of the time, incorrect pairs of sentence are chosen

Why did the chicken cross the road? The climate crisis is driven by burning fossil fuels.

Next-sentence prediction (NSP)

Next-sentence prediction

• BERT introduces two special tokens:
  • [CLS]
  • [SEP]
• We encode the input as follows \[ \mathrm{[CLS]} \bullet \text{first sentence }\bullet \mathrm{[SEP]}\bullet \text{second sentence }\bullet \mathrm{[SEP]} \]

Next-sentence prediction

• As discussed previously, if we only provide the tokens for \(\mathrm{[CLS]} \bullet \text{first sentence }\bullet \mathrm{[SEP]}\bullet \text{second sentence }\bullet \mathrm{[SEP]}\) as input to BERT
  • BERT would not be able to know that dog is from the first and play is from the second sentence
  • The segment embeddings \(\mathbf{E}_A\) and \(\mathbf{E}_B\) encode whether a token belongs to sentence A or B.

NSP

• Pass output for first token (corresponding to [CLS]) to linear layer – with only two output features, next and notNext, and pass through softmax.
• calculate the cross-entropy loss, perform backprop, update weights

NSP

• The CLS token interacts with every other token (no mask) – it “captures” information about all other tokens as needed for the classification task

The [SEP] Token and Sentence Pairs

• Some tasks require two input sentences at once (e.g., “do these two sentences mean the same thing?”).
• BERT handles this with a [SEP] token separating the two sentences, plus a learned segment embedding added to each token indicating which sentence it belongs to:

\[ \text{[CLS]} \ \text{Sentence A} \ \text{[SEP]} \ \text{Sentence B} \ \text{[SEP]} \]

• Segment embeddings are added elementwise to the token + positional embeddings, the same way positional encodings were added in Part 1 — just one more term in the sum.
• This sentence-pair format is what BERT’s own pretraining uses, and it’s directly relevant to Section 3, where we’ll see why comparing two full sentences this way becomes a computational bottleneck.

MLM and NSP

• For every training example, BERT builds one input (a sentence pair with some tokens replaced by [MASK]) and runs it through a single forward pass.

• That one pass produces both:

  • MLM predictions: output vectors at each masked position, fed into a softmax over the vocabulary to predict the original token.
  • NSP prediction: the output vector at [CLS], fed into the binary classifier to predict IsNext/NotNext.

• The two losses are then just added together \[ \mathcal{L} = \mathcal{L}_{MLM} + \mathcal{L}_{NSP} \]

• The combined loss is backpropagated through the shared BERT weights in one update step.

FFN Expansion-Contraction

BERT’s FFN Block

BERT’s FFN Block

• Why expand 4×? The attention sub-layer only mixes information that’s already there — it computes weighted averages of existing token vectors, which is a linear-ish operation with no added representational capacity. The FFN is where the model gets nonlinear capacity: projecting up to a much higher-dimensional space before the GELU gives that nonlinearity more “room” to carve out complex feature combinations, and the 4× ratio (\(3072 = 4 \times 768\)) is the ratio Vaswani et al. (2017) fixed in the original Transformer — BERT simply inherits it.
• The FFN is applied identically and independently at every position (same \(W_1, b_1, W_2, b_2\) for every token) — it’s often described as two \(1\times1\) convolutions, and it’s the main place per-token nonlinear processing happens, complementing attention’s role of moving information between positions.
• In parameter terms, the FFN dominates the model: for BERT-Base, \(2 \times 768 \times 3072 \approx 4.7M\) parameters per block just for \(W_1\) and \(W_2\) combined — roughly two-thirds of each encoder block’s parameters live in the FFN, not the attention heads.

Causal LM vs. Masked LM (recap)

Fine-tuning BERT

• Fine-tuning: Adapting the pretrained weights through additional training for a particular task.
• The pretraining is self-supervised but the fine-tuning is supervised, i.e., we need to have human-curated labels
• Pretraining BERT typically might involve billions of examples. Fine-tuning might involve a few thousand labeled examples (six orders of magnitude difference!).
• democratization: No individual or university research group has the resources to train a model such as BERT, but are able to develop sophisticated ML models by fine-tuning
• The same pre-trained BERT model can be used for various fine tuning tasks

Fine-tuning BERT for text classification

• Next-sentence prediction is an example of binary text classification (isNext, notNext)
• The following example corresponds to the binary classification explanation above.
• In general, we can classify text using more categories.

Text-classification (Python)

• We will demonstrate a Python script to classify text as PHISHING or SAFE
• Setup:

python -m pip install datasets, transformers, evaluate
python -m pip install "accelerate>=1.1.0"

Import libraries (datasets, evaluate, and transformers are from huggingface)

from datasets import DatasetDict, Dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer
import evaluate
import numpy as np
from transformers import DataCollatorWithPadding

• Get datasets (3000 examples with training: 70%; 15%: testing; 15%: validation )

from datasets import load_dataset
# available on huggingface dataset hub
dataset_dict = load_dataset("shawhin/phishing-site-classification") 

google-bert/bert-base-uncased

• This is the Hugging Face Hub identifier for the original BERT-base model Google released alongside the 2018 paper.
• The “base” (smaller) model with 12 transformer encoder blocks, hidden size 768, 12 attention heads, FFN inner dimension 3072
• uncased: text is lowercased and accent-stripped before tokenization

model_path = "google-bert/bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_path)

id2label = { 0: "Safe", 1: "Phishing"}
label2id = { "Safe": 0, "Phishing": 1}
model = AutoModelForSequenceClassification.from_pretrained(model_path,
                                                           num_labels=2,
                                                           id2label=id2label,
                                                           label2id=label2id,
                                                           attn_implementation="eager")

• The AutoModelForSequenceClassification class adds a binary classification head to the BERT base model.
• attn_implementation="eager" is needed to avoid the fused SDPA kernel, which is not available from PyTorch’s MPS (Apple GPU) backend¹.

Parameters

• By default, all of the model’s parameters (ca. 110,000,000 in this case) are trainable.
• Fine-tuning all of them is computationally expensive and, with a small dataset, risks overfitting.
• To reduce cost and (likely) improve generalization, we freeze the entire pretrained encoder (all 12 transformer blocks) and train only:
  • the pooler — a single dense layer + tanh applied to the [CLS] token, still part of base_model, which we explicitly unfreeze, and
  • the classification head added by AutoModelForSequenceClassification — untouched by the loop below since it lives outside base_model, and trainable simply because it’s newly initialized.
• Net effect: only 4 parameter tensors update during training (pooler weight/bias, classifier weight/bias)

for name, param in model.base_model.named_parameters():
    if "pooler" in name:
        param.requires_grad = True
    else:
        param.requires_grad = False

Preprocess input data

We set up a tokenizer and tokenize the input texts.

def preprocess_function(examples):
    return tokenizer(examples["text"], truncation=True)

tokenized_data = dataset_dict.map(preprocess_function, batched=True)

• The Data collator ensures all the examples in a batch have the same length

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

Evaluation

accuracy = evaluate.load("accuracy")
auc_score = evaluate.load("roc_auc")

def compute_metrics(eval_pred):
    # logits and ground truth labels
    predictions, labels = eval_pred
    # softmax
    probabilities = np.exp(predictions)/np.exp(predictions).sum(-1, keepdims=True)
    positive_class_probs = probabilities[:,1]
    auc = np.round(auc_score.compute(prediction_scores=positive_class_probs, references=labels)['roc_auc'], 3)
    predicted_classes = np.argmax(predictions, axis=1)
    acc = np.round(accuracy.compute(predictions=predicted_classes, references=labels)['accuracy'], 3)
    return {"Accuracy": acc, "AUC": auc}

• standard way to evaluate accuracy

Training parameters

lr = 2e-4
batch_size = 8
num_epochs = 10

training_args = TrainingArguments(
    output_dir="bert_phishing-classifier_teacher",
    learning_rate=lr,
    per_device_eval_batch_size=batch_size,
    per_device_train_batch_size=batch_size,
    num_train_epochs=num_epochs,
    logging_strategy="epoch",
    eval_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True
)

Fine-tune model

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_data["train"],
    eval_dataset=tokenized_data["test"],
    data_collator=data_collator,
    compute_metrics=compute_metrics
)

trainer.train()

Results

• The results on the independent validation set were: Accuracy: 88.7%, AUC94.6%

Single-sentence classification, multiple categories

• HOMEWORK You will implement a multiple-category classifier
• Sentences are given one of multiple labels

from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer
import evaluate
import numpy as np

dataset_name = "ag_news"
dataset = load_dataset(dataset_name)
# Similar code
# To get multiple classes
n_classes = dataset["train].features["label"].num_classes
model = AutoModelForSequenceClassification.from_pretrained(model_path,
                                                           num_labels=n_classes,
                                                           attn_implementation="eager")

• Four classes: 1-World, 2-Sports, 3-Business, 4-Sci/Tech

Named entity recognition

• Named-entity recognition (NER) seeks to locate and classify named entities mentioned in unstructured text into pre-defined categories such as

NER (Token Classification)

• every token position’s final hidden state \(\mathbf{h}_i \in \mathbb{R}^d\) gets its own prediction, using the same shared linear layer applied independently at each position:

\[ P(\text{label}_i) = \text{softmax}(\mathbf{h}_i W_{\text{tok}} + b_{\text{tok}}), \qquad \text{for each token } i \]

• \(W_{\text{tok}} \in \mathbb{R}^{d\times L}\), where \(L\) is the number of entity label types.
• Unlike sentence classification, [CLS] and [SEP] are typically excluded from the loss — they’re structural tokens, not real words with entity labels.
• Part of speech tagging is analogous
• Chunking: Find “chunks” of tokens that belong to the same entity (e.g., “Renal insufficiency”)

NER in Python

• For the most part, the Python code for NER is very similar to that for single-sentence classification

We need to tell BERT about the following classification

|   ID | Tag    | Description                |
|-----:|:-------|:---------------------------|
|    0 | O      | Outside of a named entity  |
|    1 | B-PER  | Beginning of a PER entity  |
|    2 | I-PER  | Inside of a PER entity     |
|    3 | B-ORG  | Beginning of a ORG entity  |
|    4 | I-ORG  | Inside of a ORG entity     |
|    5 | B-LOC  | Beginning of a LOC entity  |
|    6 | I-LOC  | Inside of a LOC entity     |
|    7 | B-MISC | Beginning of a MISC entity |
|    8 | I-MISC | Inside of a MISC entity    |

• We define the model to perform token classification rather than sentence classification

model = AutoModelForTokenClassification.from_pretrained(model_path, num_labels=9)

• Most of the rest of the code is the same except for book-keeping details that we will not review here

NER: Results

• After training the model for about 10 minutes on my laptop, it is already pretty good

from transformers import pipeline

nlp = pipeline("ner", model=model_fine_tuned, tokenizer=tokenizer)
example ="Konrad Adenauer was the first Chancellor of Germany"
ner_results = nlp(example)
print(ner_results)

• NER predictions:

Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks

• To detemrine the similarity between sentences, BERT requires that both sentences are fed into the network
• Finding the most similar pair in a collection of 10,000 sentences requires about 50 million inference computations (~65 hours) with BERT.
• The construction of BERT makes it unsuitable for semantic similarity search as well as for unsupervised tasks like clustering

Sentence-BERT

• With standard BERT, two sentences are passed to the transformer network together, and a single target value (e.g. similarity score) is predicted from that joint pass
• This setup doesn’t scale to tasks that need many pairwise comparisons — e.g. clustering a collection of sentences, which requires a similarity score between every pair
• For \(n = 10{,}000\) sentences, a full pairwise similarity matrix requires \[ n(n-1)/2 = 49{,}995{,}000 \] separate inference passes through BERT — each one a full forward pass of the transformer, since BERT has no way to compute a similarity score without processing both sentences jointly

SBERT sentence embeddings

• Creation of sentence vector: The linear classification layer of BERT is not used
• The attention outputs are pooled to derive a fixed sentence embedding

SBERT sentence embeddings

• SBERT assesses the similarity of two sentences by the cosine similarity of their vectors

\[ \mathrm{sim}(\mathbf{x},\mathbf{y}) = \frac{\mathbf{x}\cdot \mathbf{y}}{\Vert \mathbf{x} \Vert \Vert \mathbf{y} \Vert} = \frac{\sum_i A_i\cdot B_i}{\sqrt{\sum_i A_i^2}\cdot \sqrt{\sum_i B_i^2}} \]

• But nobody told BERT that the sentence embeddings it produces should be compatible with cosine similarity
• How do we teach BERT to do this?

Classification: natural language inference (NLI)

• SBERT is first trained on natural language inference (NLI) data: SNLI (570,000 pairs) + MultiNLI, each pair labeled entailment, contradiction, or neutral
• Each pair has a premise and a hypothesis, encoded independently by the siamese BERT, then mean-pooled to give sentence vectors \(u\) and \(v\)

Classification: natural language inference (NLI)

• The sentences (\(u\)=premise, \(v\)=hypothesis) are combined into a single feature vector:

\[ (u, v, |u-v|) \]

• \(|u - v|\) is the element-wise absolute difference — it makes the relationship between the two embeddings explicit, rather than making the classifier re-derive it from \(u\) and \(v\) alone
• This concatenated vector (size \(3n\), e.g. \(3\times768\)) is passed through a trainable linear layer \(W_t\) and softmax:

\[ o = \mathrm{softmax}\big(W_t\,(u, v, |u-v|)\big) \]

Classification: natural language inference (NLI)

• Trained with standard cross-entropy loss against the gold label (entailment / contradiction / neutral)
• This forces the model to learn fine-grained semantic distinctions — not just “same topic,” but “does this logically follow?”
• \(W_t\) is a training-time-only scaffold: at inference, it’s discarded — only the sentence encoder + pooling is kept

Regression Objective Function

• Now we want to train the model to recognize “similar” sentences.
• The authors use a Semantic Textual Similarity (STS) Benchmark

Regression Objective Function

• Mean squared-error (MSE) loss for the cosine similarity is used as the objective function.

• The STS Benchmark labels are on a \(y\in 0–5\) scale, so before computing the loss, the gold score \(y\) is rescaled as \(y^{\prime} = \frac{y}{2.5} - 1\). For instance

  • if \(y=5\), \(y^{\prime} = \frac{5}{2.5} - 1 = 1\).
  • if \(y=0\), \(y^{\prime} = \frac{0}{2.5} - 1 = -1\).

Regression Objective Function

• directly using the output of BERT leads to rather poor performances.
• Averaging the BERT embeddings achieves an average correlation of only 54.81, and using the CLStoken output only achieves an average correlation of 29.19.
• The SBERT/SRoBERTa approaches generally showed superior performance

Triplet Loss

• Requires triplets: an anchor \(a\), a positive \(p\) (similar to \(a\)), and a negative \(n\) (dissimilar to \(a\)) — no graded similarity score needed, just relative ordering

• We assume that a sentence is thematically closer to sentences within its section than to sentences from other sections

• \(a\): McDonnell resigned from Martin in 1938 and founded McDonnell Aircraft Corporation in 1939

• \(p\): In 1967, McDonnell Aircraft merged with the Douglas Aircraft Company to create McDonnell Douglas

• \(n\): Born in Denver, Colorado, McDonnell was raised in Little Rock, Arkansas, and graduated from Little Rock High School in 1917

Triplet Loss

• Compute two Euclidean distances: \[ d_{ap} = \|s_a - s_p\|, \qquad d_{an} = \|s_a - s_n\| \]

• Loss (with margin \(\epsilon\)): \[ \mathcal{L} = \max\big(d_{ap} - d_{an} + \epsilon,\ 0\big) \]

• The margin \(\epsilon\) ensures that \(s_p\) is at least \(\epsilon\) closer to \(s_a\) than \(s_n\).

• As metric Euclidean distance was used.

• \(\epsilon=1\)

Triplet Loss — role of the margin \(\epsilon\)

• If \(p\) is already closer to \(a\) than \(n\) is, by at least \(\epsilon\) → loss is 0, no gradient, nothing to learn
• Otherwise → loss is positive, and gradients pull \(a\) and \(p\) together / push \(a\) and \(n\) apart until the margin is satisfied

• Without the margin, the loss would be 0 even if \(d_{ap}\) is only slightly smaller than \(d_{an}\) — a weak signal that stops training too early
• The margin forces a gap, producing embeddings with clearer separation between similar and dissimilar sentences

SBERT: three objectives

SBERT - Summary

• SBERT maps sentences to a vector space that better reflects sentence emantic similarity than BERT’s sentence vectors
• SBERT is computationally efficient. SBERT can be used for tasks which are computationally not feasible to be modeled with BERT (e.g., clustering of large corpora of sentences).
• We will see in a later lecture that SBERT is often used to implement Retrieval Augmented Generation (RAG)

Sources

• Devlin J et al. (2019) BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. arXiv 1810.04805
• Hugging Face (blog) BERT 101 🤗 State Of The Art
• Rogers A et al. (2020) A Primer in BERTology: What we know about how BERT works. arXiv 2002.12327
• Reimers N and Gurevych I (2019)Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks. arXiv 1908.10084