Skip to content

Repository files navigation

Extending BERT with Multi-task and Meta-learning

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.

Paper Python 3.8 PyTorch 1.8 License

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.


Results

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%

What moved the needle

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%

Few-shot: Proto-BERT

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%.


Approach

1. minBERT from scratch

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.

2. Shared projected attention layers (PALs)

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.

3. Siamese sentence encoders

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.

4. Proto-BERT for few-shot sentiment

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.


Quickstart

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_dfp

Or install into an existing environment:

pip install -r requirements.txt

All 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 64

Multi-task model (all three tasks, shared encoder):

bash run.sh

run.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/

Verifying the implementation

python -m tests.sanity_check     # checks BERT outputs against reference embeddings
python -m tests.optimizer_test   # checks AdamW against reference parameter values

Both should report success. sanity_check downloads bert-base-uncased weights on first run.

Packaging a submission

python -m prepare.prepare_submit

Repository layout

models/       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.

Scope of this repository

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.


Datasets

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

Error analysis

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?"


Citation

@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}
}

Acknowledgements

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).

License

Apache License 2.0 — see LICENSE.

About

Extending BERT with multi-task learning and meta-learning — shared projected attention layers, Siamese sentence encoders, and Proto-BERT few-shot classification. Top 5 on the Stanford CS 224N leaderboard, 74.0% average accuracy (+20.4% over baseline)

Topics

Resources

Stars

1 star

Watchers

1 watching

Forks

Releases

Packages

Contributors

Languages