peterkirby/tinybert-pan2020-authorship-verification

TinyBERT PAN Authorship Verification

A full fine-tune of huawei-noah/TinyBERT_General_4L_312D for comparing whether two English fanfiction documents share an author. Training used peterkirby/pan2020_dict_author_fandom_doc, a contrastive batch of 1,024 documents, and a maximum input length of 512 tokens per document.

This is the 40-epoch, mean-pooling-only TinyBERT model from Table 4 of Contrastive Learning for Authorship Verification. Its selected checkpoint achieved 86.1400% validation accuracy and 0.8627315164 validation binary F1. The saved validation-selected rule achieved 86.5843% accuracy on PAN21 test.

The larger companion is ModernBERT-large PAN Authorship Verification, which reports 98.4199% PAN21 test accuracy with 4,096-token inputs. Both models use masked mean pooling, normalized embeddings, and a saved cosine threshold. TinyBERT has 14,350,248 encoder parameters and 312-dimensional embeddings; ModernBERT-large has 394,781,696 encoder parameters and 1,024-dimensional embeddings. Their input lengths and training configurations differ, so this is a comparison of the released systems.

Model and embedding representation

The encoder has 4 Transformer layers, 312-dimensional hidden states, and 14,350,248 encoder parameters. It uses the base TinyBERT BERT vocabulary, attention-mask-aware mean pooling over all non-padding tokens, and L2 normalization. Special tokens participate in pooling. There is no learned projection head or separate classifier head. The BERT pooler tensors are retained from the checkpoint, but document embeddings use last_hidden_state and masked mean pooling rather than the pooler output.

The export preserves all 71 encoder tensors in their original float32 dtype, without retraining or quantization. model.safetensors contains the encoder. The exact decision threshold is stored in both config.json under authorship_verification and authorship_verification_config.json.

The architecture and the reported fine-tune support 512 tokens, including special tokens. Longer texts are truncated on the right for the recorded evaluation. Chunking and averaging longer texts would define a different evaluation setup.

Binary threshold and decision rule

The original Lightning checkpoint stores state_dict["eval_threshold"] as 0.7344273328781128. For normalized document embeddings u and v, apply:

same_author = cosine_similarity(u, v) >= 0.7344273328781128
label = int(same_author)  # 1: same author; 0: different authors

At each validation epoch, the training code searched precision-recall thresholds and selected the threshold with the highest binary F1 on pan21/validation. Ties used the first maximum returned by NumPy's argmax. Lightning selected the checkpoint with the highest val_f1; the saved threshold was then reused for the test splits. The recorded threshold-selection and test code is in contrastive-av at the training commit.

Equality belongs to the same-author class. This release applies the saved binary rule directly, with no sigmoid, probability calibration, abstention band, or score remapping. The cosine score is a similarity score, not a probability.

Dataset and evaluation

The training loader uses the dataset's default configuration and train split, with author IDs as supervised contrastive labels. The dataset documents an author-and-fandom-disjoint validation split. Threshold selection uses pan21/validation; testing uses pan21/test and pan20/test. See the dataset card for preprocessing, split construction, and attribution.

Evaluation split	Accuracy	Binary F1	Threshold
PAN21 validation at selected checkpoint	86.1400%	0.8627315164	0.7344273328781128
PAN21 test	86.5843%	0.8613365293	0.7344273328781128
PAN20 test	86.1645%	0.8653610945	0.7344273328781128

These are the original run's results. evaluation_results.json stores accuracy and F1 as fractions from 0 to 1; the table expresses accuracy as percentages. training_history.parquet contains all 385 original history records, including the 40 validation epochs. The selected checkpoint is epoch 39, the 40th epoch. The full benchmark was not rerun during export.

The original trainer used bf16-mixed precision and FlashAttention 2; the saved encoder weights are float32. Evaluation used float32 pooling and cosine scores, 512-token inputs, right padding and truncation, and batches of 256 documents. Each side of a document pair was encoded separately. Alternate precision, attention implementations, or padding can affect scores near the threshold.

Training details

Setting	Value
Base encoder	huawei-noah/TinyBERT_General_4L_312D
Objective	Supervised contrastive loss (SupConLoss)
Temperature	0.01
Contrastive batch	1,024 documents on one GPU
Sampler	M-per-class, 2 documents per author
Training span sampling	Random character/word-boundary rotation before tokenization and truncation
Evaluation span sampling	No random rotation; prefix of each document
Maximum sequence length	512 tokens, including special tokens
Gradient-cache microbatch	256 documents
Optimizer	bitsandbytes PagedAdamW8bit
Learning rate	0.0004
Weight decay	0.01
Adam betas / epsilon	(0.9, 0.999) / 0.00000001
Schedule / warmup	Linear / 10%
Training precision	bf16-mixed
Seed	1000
Hardware	1 × NVIDIA RTX 4090
Training budget	40 epochs
Selected checkpoint	Epoch 39 (zero-based; 40th epoch), global step 15,200
Selection metric	Maximum validation binary F1
Recorded training Git commit	aa8671590e47cc364be005dcb8edb60cbcd01474

training_config.yaml is the resolved configuration saved by the original run, including the head_type=none and pooling=mean overrides. The repository file named tinybert_contrastive_best.yaml had different pooling/head defaults; the resolved run configuration identifies this checkpoint. training_hyperparameters.json comes directly from the checkpoint. training_requirements.txt preserves the original dependency inventory.

The training configuration did not pin base-model or dataset revisions. provenance.json records the revisions used to fetch tokenizer/config assets and inspect reference cards during export; those are not asserted to be the training-time revisions.

Use with Transformers

Install requirements.txt for the pinned inference dependencies. This example uses PyTorch SDPA, with float32 on CPU and bfloat16 on CUDA:

import numpy as np
import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer

model_id = "peterkirby/tinybert-pan2020-authorship-verification"
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if device == "cuda" else torch.float32
tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side="right")
model = AutoModel.from_pretrained(
    model_id, dtype=dtype, attn_implementation="sdpa",
).to(device).eval()
settings = model.config.authorship_verification

@torch.inference_mode()
def embed(text):
    inputs = tokenizer(
        [text], add_special_tokens=True, padding="max_length", truncation=True,
        max_length=settings["max_length"], return_tensors="pt",
    ).to(device)
    hidden = model(**inputs).last_hidden_state
    mask = inputs.attention_mask.unsqueeze(-1).expand(hidden.size()).float()
    pooled = (hidden * mask).sum(1) / mask.sum(1).clamp(min=1e-9)
    return F.normalize(pooled, p=2, dim=-1)

first = embed("First document text...")
second = embed("Second document text...")
similarity = F.cosine_similarity(first, second).item()
same_author = similarity >= settings["threshold"]
print({"cosine_similarity": similarity, "same_author": same_author})

AutoModel returns token states. The mean pooling and normalization are necessary to obtain the document embedding. The generic feature-extraction pipeline alone does not apply this pooling or the authorship decision rule.

The bundled authorship_verification.py provides batched embeddings and the binary rule. From a downloaded copy of this repository:

python authorship_verification.py --model . \
  --text1-file /path/to/document1.txt \
  --text2-file /path/to/document2.txt

For the original evaluation attention implementation, use a compatible FlashAttention 2 build and select --device cuda --attn-implementation flash_attention_2. Local inference can use this directory as the model path. No custom Transformers remote-code loading is required.

Release files and verification

File	Purpose
model.safetensors	Fine-tuned encoder tensors in float32
config.json, tokenizer files	Architecture, tokenizer, and inference settings
authorship_verification_config.json	Pooling, sequence length, labels, and exact threshold
authorship_verification.py	Document embedding and binary comparison helper
evaluation_results.json, training_history.parquet	Original evaluation results and training history
export_verification.json	Tensor, tokenizer, pooling, and decision-boundary checks
provenance.json, training_*	Checkpoint identity and training configuration
checkpoint/model.ckpt	Original Lightning training checkpoint
BASE_MODEL_README.md	Base model's original model card and attribution
SHA256SUMS	File checksums for this release

Export verification confirmed that all 71 encoder tensors and their dtypes match the checkpoint exactly, tokenization matches the fetched base assets at 512 tokens, and the bundled helper reproduces the source encoder plus training-time pooling exactly on a CPU float32 SDPA smoke test. The smoke inputs include empty text and a long input requiring truncation. Embeddings are finite, L2-normalized, and 312-dimensional. Scores immediately below, equal to, and above the threshold produce labels 0, 1, and 1.

The export was verified with PyTorch 2.10.0 and Transformers 5.3.0. The original dependency inventory records PyTorch 2.10.0+cu128, Transformers 5.6.2, and FlashAttention 2.8.3. These checks validate the export and decision rule; they do not establish a new benchmark result or numerical equivalence with the original CUDA mixed-precision execution.

To verify a downloaded release locally:

sha256sum -c SHA256SUMS

Original Lightning checkpoint

checkpoint/model.ckpt is the unmodified 86,748,001-byte checkpoint, originally named epoch=39-step=15200-val_f1=0.8627.ckpt. It includes the encoder, saved threshold, optimizer state, scheduler state, and training hyperparameters. Its SHA-256 checksum is:

49d4a8d6ef67a6d9cb6923d06db29ca6f5bce8cd679e00a3a120f6291a2aae55

The source is W&B run tbk93d2k, artifact peterkirby/contrastive-av/model-tbk93d2k:v0. The checkpoint uses model._orig_mod.encoder.* keys from torch.compile; the Transformers export removes that prefix. Restoring Lightning training requires the original classes and matching environment. AutoModel.from_pretrained loads the safetensors export instead.

Intended use and limits

Use this model for research and evaluation of pairwise authorship verification on English fanfiction. The recorded results characterize the named PAN splits. Performance on other genres, languages, short snippets, edited text, or a different author distribution has not been established by this release. The model compares textual style; it does not identify a real person or establish authorship with certainty.

For a larger encoder and longer evaluated context, see ModernBERT-large PAN Authorship Verification.

Citation

If you use this model or its authorship-verification method, please cite:

Peter Kirby. Contrastive Learning for Authorship Verification. In Experimental IR Meets Multilinguality, Multimodality, and Interaction (CLEF 2026), Lecture Notes in Computer Science, volume 17087, pages 92–102. Springer, 2027.

@inproceedings{kirby2027contrastive,
  author = {Kirby, Peter},
  title = {Contrastive Learning for Authorship Verification},
  booktitle = {Experimental IR Meets Multilinguality, Multimodality, and Interaction},
  series = {Lecture Notes in Computer Science},
  volume = {17087},
  pages = {92--102},
  year = {2027},
  publisher = {Springer},
  doi = {10.1007/978-3-032-39150-6_7},
  url = {https://doi.org/10.1007/978-3-032-39150-6_7}
}

An open preprint is available on arXiv. For the base architecture, see TinyBERT: Distilling BERT for Natural Language Understanding and BASE_MODEL_README.md.

Source licensing and attribution

The fetched base-model Hub repository contains no explicit license metadata or license file; its original model card is preserved as BASE_MODEL_README.md. The training dataset declares the custom pan-2020-fanfiction-dataset license. No separate fine-tuned-weight license was recorded in the training artifact, so this card leaves that metadata unspecified. Follow the base-model and dataset sources for attribution and terms.