A from-scratch BERT implementation extended with multi-task learning, Siamese sentence encoders, and few-shot meta-learning — finishing in the top 5 of the Stanford CS 224N class leaderboard.
Jing Ning, Cam Burton — Stanford University · CS 224N Final Project
📄 Read the paper (Stanford CS 224N final reports archive)
No
transformers, no pretrained model wrappers. The transformer encoder, multi-head self-attention, and the AdamW optimizer are all implemented from scratch in this repo and verified against reference outputs.
A single shared BERT encoder serves three tasks at once. Final held-out test accuracy of the ensemble model:
| Task | Dataset | Metric | Score |
|---|---|---|---|
| Sentiment analysis | SST-5 | Accuracy | 54.2% |
| Paraphrase detection | Quora | Accuracy | 81.0% |
| Semantic textual similarity | SemEval STS | Accuracy | 86.8% |
| Average | 74.0% |
+20.4% average accuracy over the fine-tuned multi-task baseline.
Development-set scores for the three ensemble members and their combination:
| Model | Sentiment | Paraphrase | Semantic similarity | Average |
|---|---|---|---|---|
| model 1 | 54.9% | 81.8% | 86.4% | 74.1% |
| model 2 | 53.2% | 82.2% | 86.1% | 73.6% |
| model 3 | 53.3% | 81.4% | 86.8% | 73.6% |
| Ensemble | 55.8% | 83.2% | 88.9% | 75.9% |
Each component measured against its own experiment's baseline (deltas are not strictly additive):
| Component | Δ Average accuracy |
|---|---|
| Additional out-of-domain training data | +8.6% |
| Siamese u−v interaction + 3-layer LSTM heads | +7.3% |
| Shared projected attention layer (PALs) | +4.2% |
| MSE regression loss for similarity | +3.1% |
| Ensembling three models | +2.4% |
Mean-of-hidden-layers output (vs. [CLS]) |
+1.7% |
| Multi-task loss weighting (1.5× paraphrase) | +0.6% |
| Linear–cosine–exponential LR schedule | +0.2% |
| Overall vs. multi-task baseline | +20.4% |
A prototypical-network formulation of BERT for 5-way 5-shot sentiment classification, evaluated out of domain on Amazon Kindle book reviews after meta-training on SST:
| Meta-test set | 5-way 5-shot query accuracy |
|---|---|
| Amazon Kindle reviews (out of domain) | 38.7% ± 5.7% |
| SST dev (in domain) | 41.7% ± 21.1% |
Only five labelled examples per class — and the out-of-domain gap is just 3%.
models/bert.py implements the encoder end to end: scaled dot-product multi-head
self-attention, the add-norm residual blocks, and the embedding layer (token + positional +
segment). training/optimizer.py implements AdamW with decoupled weight decay and the
efficient bias-correction form. Both are checked against reference outputs — see Verifying the implementation.
Following Stickland & Murray, a single low-dimensional attention layer is shared across all three task heads. Sharing one attention layer beat fully independent heads by +4.2% average accuracy, and paraphrase accuracy alone by +12% — the shared layer learns cross-task features while each head specialises.
Both pair tasks encode each sentence independently and feed the head a concatenation of
[u, v, u − v]. Two changes to the Sentence-BERT recipe:
introducing the u − v difference term early, and replacing the pooling layer with three
stacked LSTMs (one each for u, v, and u − v). Together, +7.3% over the pooling
baseline. A 1D-CNN feature extractor was also tried and came out −4.5% behind.
Class prototypes are the mean BERT embedding of each class's support set; queries are classified by softmax over negative Euclidean distance to each prototype, with an episodic data loader for meta-training.
Full method and ablations are in the paper.
git clone https://github.com/CDC1688/BERT-multitask-metalearning.git
cd BERT-multitask-metalearning
bash setup.sh # creates the cs224n_dfp conda env
conda activate cs224n_dfpOr install into an existing environment:
pip install -r requirements.txtAll commands run from the repository root.
Sentiment baseline (SST-5 and CFIMDB, single task):
python -m training.classifier --option finetune --use_gpu --epochs 10 --lr 1e-5 --batch_size 64Multi-task model (all three tasks, shared encoder):
bash run.shrun.sh is parameterised by environment variables:
EPOCHS=30 LR=3e-5 BATCH_SIZE=64 RUN_NAME=my-run bash run.sh--option pretrain freezes BERT and trains only the heads; --option finetune updates the
full encoder. Predictions are written to predictions/, TensorBoard logs to runs/:
tensorboard --logdir runs/python -m tests.sanity_check # checks BERT outputs against reference embeddings
python -m tests.optimizer_test # checks AdamW against reference parameter valuesBoth should report success. sanity_check downloads bert-base-uncased weights on first run.
python -m prepare.prepare_submitmodels/ BERT architecture
bert.py ★ Multi-head self-attention, add-norm, encoder
base_bert.py Weight loading and initialisation (provided)
tokenizer.py WordPiece tokenizer (provided)
config.py Model configuration (provided)
utils.py Download and caching helpers (provided)
training/ Training entry points and optimisation
multitask_classifier.py ★ Shared encoder + three task heads
classifier.py ★ Single-task sentiment baseline
optimizer.py ★ AdamW with decoupled weight decay
schedulers.py ★ Linear / cosine / exponential LR schedules
checkpoint.py ★ Checkpoint save and load
evaluation/ Metrics and prediction writing
metrics.py Per-task scoring, dev and test (provided)
prepare/ Data loading and packaging
datasets.py Dataset classes and CSV loading (provided)
prepare_submit.py Submission zip builder
tests/ Reference-output tests
sanity_check.py BERT encoder vs. reference embeddings
optimizer_test.py AdamW vs. reference parameter values
data/ SST, CFIMDB, Quora, and STS splits
predictions/ Model outputs (written at test time)
★ = implemented for this project, as opposed to course-provided scaffolding.
This repo contains the multi-task system: the from-scratch BERT encoder, the AdamW optimizer, the single-task baseline, and the shared-encoder multi-task model with Siamese pair heads and LR scheduling. The PALs layer, the stacked-LSTM heads, Proto-BERT, and the ensembling scripts were developed in separate experiment branches and are described in the paper; the numbers above are the paper's reported results, not the output of a single command in this tree.
| Dataset | Task | Labels |
|---|---|---|
| SST | Sentiment | 5-way, negative → positive |
| CFIMDB | Sentiment | Binary |
| Quora | Paraphrase detection | Binary |
| SemEval STS | Semantic similarity | 0–5 continuous |
| SICK | Similarity (supplementary) | 0–5 continuous |
| Amazon Kindle reviews | Few-shot meta-test | 5-way |
Sentiment. Per-class accuracy varies by as much as 32% across the five classes. Of the incorrect predictions, roughly 91.5–95% fall in an adjacent sentiment class — the model learns the ordinal structure of sentiment, but the boundaries between neighbouring categories are genuinely blurry. Training accuracy above 90% against 54% test accuracy shows overfitting remains the binding constraint on this task.
Paraphrase. Non-duplicate pairs are classified 9.6% more accurately than duplicates. Manual review of misclassified examples suggests some labels are debatable for human annotators too:
"What are some interesting campus recruitment rejection stories?" "What are some of the best rejection stories at campus recruitment?"
@techreport{ning2023extending,
title = {Extending BERT with Multi-task and Meta-learning},
author = {Ning, Jing and Burton, Cam},
institution = {Stanford University},
type = {CS 224N Final Report},
year = {2023},
url = {https://web.stanford.edu/class/archive/cs/cs224n/cs224n.1234/final-reports/final-report-169919951.pdf}
}The BERT scaffold is adapted from the minbert assignment developed for Carnegie Mellon University's CS11-711 Advanced NLP by Shuyan Zhou, Zhengbao Jiang, Ritam Dutt, Brendon Boldt, Aditya Veerubhotla, and Graham Neubig, and from the Stanford CS 224N default final project.
Parts of the code derive from the
transformers library
(Apache License 2.0).
Apache License 2.0 — see LICENSE.