ESM-2 Transformer Head For Full EC Multi-Label Prediction

Frozen ESM-2 residue representations and a trainable Transformer head are used to predict complete four-part enzyme commission labels such as 2.7.11.1. Unlike EC level-1 classification, this experiment distinguishes 5,320 EC labels and supports more than one correct EC number per protein.

Final held-out test result: full-label micro F1 0.9344, macro F1 0.7501, sample F1 0.9505, and exact match 0.9191.

Contents

• Task
• Model architecture
• Data preparation
• Leakage and unseen-label control
• Hyperparameter and loss selection
• Training and early stopping
• Final results
• Reproducing the experiment
• Repository contents
• Security and excluded artifacts
• Limitations
• License

Task

Each input is a protein sequence. A sequence can have one or several complete EC annotations.

Annotation string	Meaning
1.1.1.1	One positive full EC label
1.1.1.1;2.7.11.1	Two independent positive full EC labels

This is a multi-label problem:

• the classifier produces 5,320 independent logits;
• sigmoid is applied independently to each output;
• training uses BCEWithLogitsLoss;
• softmax and CrossEntropyLoss are not used;
• validation-selected thresholds convert scores into predicted EC sets.

Model Architecture

flowchart LR
    A[Protein sequence] --> B[ESM-2 tokenizer]
    B --> C["Frozen ESM-2<br/>t33 650M UR50D"]
    C --> D["Residue token embeddings<br/>1280 dimensions"]
    D --> E["Linear projection<br/>1280 to 256"]
    E --> F["Transformer encoder head<br/>1 layer, 8 heads"]
    F --> G[Masked mean pooling]
    G --> H["Linear classifier<br/>256 to 5320"]
    H --> I[5,320 independent logits]

Loading

Component	Configuration	Trainable?
ESM-2 encoder	facebook/esm2_t33_650M_UR50D	No
ESM-2 hidden size	1,280	No
Input projection	1280 -> 256	Yes
Transformer layers	1	Yes
Attention heads	8	Yes
Feed-forward dimension	512	Yes
Activation	GELU	Yes
Dropout	0.1	Yes
Pooling	Masked residue-token mean	No parameters
Classifier	256 -> 5320	Yes

The final model receives token-level ESM-2 representations. It does not feed one already mean-pooled ESM-2 vector into the Transformer.

Data Preparation

Raw input

The raw source is not redistributed. The local experiment used a CSV-formatted file with these columns:

Column	Description
id	Protein identifier
ec_number	Complete EC annotation, optionally followed by evidence text
seq	Amino-acid sequence

Audit item	Value
Raw annotation rows	291,538
Raw unique IDs	238,346
Accepted complete numeric EC rows	290,119
Excluded provisional or invalid rows	1,419
Raw SHA256	fcd14e264f78c70652c59db812c819fbdc78551626503a7772b8003408ee0d8b

Only numeric four-part labels matching 1-7.NUMBER.NUMBER.NUMBER were used. Evidence suffixes such as {ECO:...} were removed after EC extraction. Provisional labels such as 6.2.1.n2 were excluded and counted in the audit. The raw file was never overwritten.

Processing flow

flowchart TD
    A[Raw id / EC / sequence rows] --> B[Extract complete numeric four-part EC]
    B --> C[Group annotations by locked sequence ID]
    C --> D[Union and numerically sort EC labels]
    D --> E[Reuse locked train / validation / test IDs]
    E --> F[Build vocabulary from train only]
    F --> G[Save processed split CSVs and audit JSON]

Loading

The previous locked sequence split was reused rather than making a favorable new split for the full EC task.

Split	Proteins	EC annotations	Unique truth labels	Mean labels/protein
Train	160,047	172,382	5,320	1.0771
Validation	20,006	21,572	2,607	1.0783
Test	20,006	21,550	2,563	1.0772

Leakage And Unseen-Label Control

Audit	Result
Train-validation ID overlap	0
Train-test ID overlap	0
Validation-test ID overlap	0
Exact sequence overlap between any split pair	0
Duplicate IDs within each split	0
Duplicate sequences within each split	0
Sequence length range	50 to 1,024 residues

The output vocabulary comes from train only. Consequently, some labels in validation and test cannot be emitted by the classifier.

Split	Unseen labels	Affected proteins	Unseen annotations
Validation	179	187	192
Test	184	180	189

Two views are reported:

• known-vocabulary metrics evaluate the 5,320 outputs available to the model;
• full-label metrics additionally count unseen truth labels as false negatives.

Full-label metrics are the primary results in this README. Unseen labels are never silently dropped from the main score.

Label Imbalance

Full EC labels are extremely sparse. Of the 5,320 train labels, 1,475 occur once and 2,900 occur fewer than five times.

Train support per label	Number of labels
1	1,475
2-4	1,425
5-9	759
10 or more	1,661

Hyperparameter And Loss Selection

The level-1 loss choice was not assumed to work for 5,320 sparse outputs. A small validation-only screening compared three loss settings while holding the model, learning rate, seed, subset, and epoch count fixed.

Fixed screening setting	Value
Train subset	8,192 proteins
Validation subset	4,096 proteins
Epochs	2
Seed	1234
Head learning rate	1e-4
Transformer	256 dimensions, 1 layer, 8 heads
Test loaded	No

Candidate	Validation full micro F1	Full macro F1	Decision
Unweighted BCE	0.0168	0.0001	Rejected
Weighted BCE, cap 20	0.0387	0.0014	Rejected
Weighted BCE, cap 100	0.5142	0.1774	Selected

With 5,320 outputs, unweighted BCE was dominated by negative output positions. The selected positive weights were computed from the complete train split only and capped at 100. The uncapped maximum would have been 160,046.

This was targeted loss screening. It was not an exhaustive search over every architecture, learning rate, pooling method, or seed.

Training And Early Stopping

Setting	Value
Optimizer	AdamW
Head learning rate	1e-4
Weight decay	0.01
Batch size	16
Mixed precision	Enabled on CUDA
Maximum epochs	15
Minimum epochs	8
Early-stopping patience	4
Checkpoint criterion	Validation known-vocabulary micro F1
Best checkpoint	Epoch 8
Stopped	Epoch 12
Recorded epoch time	Approximately 8.36 hours total

Train loss continued decreasing after epoch 8, while validation loss increased and validation micro F1 declined. This is why the epoch-8 checkpoint, rather than the final epoch-12 state, was evaluated on test.

The checkpoint and threshold procedure was:

1. Train on the train split.
2. Select the checkpoint by validation micro F1.
3. Tune global and eligible per-label thresholds on validation only.
4. Reload the selected checkpoint.
5. Evaluate the untouched test split once.

Threshold Policy

The selected epoch used validation-only thresholds:

• one global fallback threshold of 0.9;
• individual thresholds for 496 labels with at least 10 validation positives;
• global fallback for all other train-vocabulary labels;
• top-1 fallback when no label crosses its threshold.

Threshold	Number of output labels
0.05	7
0.10	24
0.20	62
0.30	59
0.40	59
0.50	43
0.60	46
0.70	47
0.80	61
0.90	4,912

Because weighted BCE changes the score distribution, sigmoid scores should not automatically be interpreted as calibrated biological probabilities. They are classification scores used with validation-selected decision thresholds.

Final Results

Primary full-label test metrics

Metric	Validation	Test
Micro F1	0.9391	0.9344
Macro F1 over truth labels	0.7600	0.7501
Sample F1	0.9534	0.9505
Exact match	0.9236	0.9191

Test decision count	Value
True positives	20,741
False positives	2,104
False negatives	809
Average predicted labels/protein	1.1419
Average true labels/protein	1.0772

Known-vocabulary diagnostic metrics

Metric	Test
Micro F1	0.9384
Macro F1 over supported labels	0.8081
Sample F1	0.9506
Known-vocabulary subset accuracy	0.9195
Known-vocabulary Hamming loss	0.0000256

The full-label micro F1 is lower because 189 test annotations belong to EC labels that never appear in train.

Frequent-label behavior

The aggregate score is strong for common labels, while the macro score remains lower because thousands of labels have very little supervision. The complete 5,320-row test class table is available at results/level4_transformer_head_cap100_s1234_test_class_metrics.csv.

Reproducing The Experiment

1. Install

Python 3.10 or newer and a CUDA-capable GPU are recommended.

python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

No OpenAI or Codex API key is needed.

2. Prepare local data

Place private inputs locally as follows:

data/
├── raw/
│   └── full_ec_annotations.csv
└── locked_splits/
    ├── cached_full_t33_lr3e4_posweight_train.csv
    ├── cached_full_t33_lr3e4_posweight_valid.csv
    └── cached_full_t33_lr3e4_posweight_test.csv

The raw file requires id, ec_number, and seq. The locked split files also require those columns; their EC level-1 column is replaced by recovered full EC annotations during preprocessing.

python src/prepare_data.py \
  --raw data/raw/full_ec_annotations.csv \
  --source_split_dir data/locked_splits \
  --output_dir data/processed

3. Optional loss screening

python src/run_loss_screen.py

This command is validation-only and does not load test.

4. Full training and test evaluation

python src/train.py \
  --train_csv data/processed/level4_train.csv \
  --valid_csv data/processed/level4_valid.csv \
  --test_csv data/processed/level4_test.csv \
  --vocabulary data/processed/train_label_vocabulary.txt \
  --output_dir outputs/full_run \
  --run_name level4_transformer_head_cap100_s1234 \
  --epochs 15 \
  --min_epochs 8 \
  --patience 4 \
  --batch_size 16 \
  --eval_batch_size 16 \
  --head_lr 1e-4 \
  --weight_decay 0.01 \
  --dropout 0.1 \
  --transformer_dim 256 \
  --transformer_layers 1 \
  --transformer_heads 8 \
  --transformer_ff_dim 512 \
  --pooling mean \
  --loss_mode pos_weight \
  --pos_weight_cap 100 \
  --threshold_min_validation_positives 10 \
  --seed 1234

5. Regenerate public figures

python src/plot_results.py

Repository Contents

.
├── assets/                       # README figures
├── results/                      # Small aggregate result artifacts
├── src/
│   ├── dataset.py                # Sparse label dataset and batch collation
│   ├── metrics.py                # Known/full metrics and threshold tuning
│   ├── model.py                  # Frozen ESM-2 Transformer-head model
│   ├── plot_results.py           # Rebuilds README figures
│   ├── prepare_data.py           # Full EC normalization and locked split join
│   ├── run_loss_screen.py        # Validation-only loss screening
│   └── train.py                  # Training, early stopping, and final test
├── vocabulary/
│   └── train_label_vocabulary.txt
├── .gitignore
├── LICENSE
├── requirements.txt
├── SECURITY.md
├── THIRD_PARTY_NOTICES.md
└── README.md

Security And Excluded Artifacts

Only code, small aggregate result files, the train-label vocabulary, and figures are intended for version control.

The following are deliberately excluded:

• OpenAI, Codex, GitHub, or other API keys and access tokens;
• .env and credential files;
• raw protein data and processed sequence split CSVs;
• complete per-protein prediction CSVs;
• ESM-2 model cache;
• PyTorch checkpoints and optimizer state;
• logs, notebook caches, and Python bytecode.

The .gitignore blocks these common paths. Always review git status and run a secret scan before publishing.

Limitations

1. This is a single-seed (1234) result.
2. Exact duplicate sequences do not cross splits, but the split is not based on sequence-identity clustering; it is not a remote-homology benchmark.
3. There are 1,475 labels with one train example and 2,900 with fewer than five.
4. Test contains 184 labels unseen in train; they are counted as false negatives.
5. ESM-2 was frozen; end-to-end fine-tuning was not tested.
6. Weighted BCE scores are not probability-calibrated.
7. Loss weighting was screened, but architecture and learning rate were not exhaustively optimized for full EC prediction.
8. The raw dataset and checkpoint are not redistributed in this repository.

Result Interpretation

The strongest supported statement is:

On the locked held-out split, the frozen ESM-2 token-level Transformer head achieved full-label micro F1 0.9344 and macro F1 0.7501 for 5,320 train-derived complete EC outputs, while counting unseen test labels as false negatives.

This should not be described as performance on an independent future dataset, a homology-clustered benchmark, or a three-seed robustness study.

License

The original code in this repository is released under the MIT License. Third-party software and model components remain under their respective licenses; see THIRD_PARTY_NOTICES.md.

The raw protein dataset, processed splits, complete prediction file, ESM-2 weights, and trained checkpoint are not covered as redistributed artifacts because they are not included in this repository.