saugataroy/ielts-task2-modernbert-base-moe

IELTS Task 2 ModernBERT-base — Multi-gate MoE

An English essay-scoring research model that estimates Task Response (TR), Coherence and Cohesion (CC), Lexical Resource (LR), and Grammatical Range and Accuracy (GRA) from an IELTS-style Task 2 question and typed response.

This checkpoint continues the ELLIPSE-to-IELTS MoE model with synthetic augmentation. Its labels come from Writing9 and LLM judgments. Reliable agreement with qualified IELTS examiners has not been established. It is not an official IELTS grading service.

Model details

Property	Value
Publisher	saugataroy
Architecture	ModernBERT-base encoder, essay-token mean pooling, dense multi-gate MoE head
Parameters	150,078,228
Language and scope	English; Writing Task 2 only
Input limit	1,024 tokens including question, response, and special tokens
Weight format	Full FP32 Safetensors; encoder and head
Parent checkpoint	ad150d954cc9af7fa848fedeafeef6c28c71e2c6
Synthetic dataset revision	067d907b930a8711a4c9888a2258ede3fb99b83e
Related plain-head model	saugataroy/ielts-task2-modernbert-base-mlp

Uses and limitations

Intended for rubric-regression research, writing-practice prototypes, and controlled studies of synthetic augmentation. Human review and independent examiner-rated evaluation are needed before relying on its estimates for consequential decisions. Do not use it to determine official examination bands, admissions, or employment.

Task 1 Academic, General Training letters, handwriting, speaking, and a complete IELTS Writing band are unsupported. The model supplies estimates, not generated feedback, sentence-level corrections, or calibrated confidence. No subgroup fairness evaluation or representative real-world examiner study has been completed.

Synthetic responses can introduce artificial errors, model-specific style, and rubric bias. Their GRA labels are strongly skewed high: 533 of 1,194 accepted rows are band 9, while only three are band 3. Synthetic holdout gains measure agreement with LLM labels. Original-corpus results measure agreement with Writing9, whose scoring reliability is not independently established. Neither result should be described as official IELTS accuracy.

Inputs and outputs

Provide the original question and a nonempty essay. Preserve spelling, grammar, and paragraphs. The paired tokenizer preserves the question and truncates only the essay when necessary; overly long questions can fail tokenization.

model.score returns one dictionary per input pair:

• scores: continuous TR, CC, LR, and GRA estimates on the training labels' 0–9 scale.
• task2_mean_estimate: their unrounded arithmetic mean.
• truncated: whether the full input exceeded 1,024 tokens.
• track: ielts; max_tokens: 1024.

Outputs are not clipped or rounded; values outside the label range are possible. The mean is not an official rounded Task 2 band or the overall Writing score.

Usage

Use Python 3.11–3.13 and a PyTorch build appropriate for your device:

pip install 'torch>=2.6,<2.12' 'transformers==4.57.6' 'huggingface-hub==0.36.0' 'safetensors>=0.6,<0.8'

from pathlib import Path
from transformers import AutoModel, AutoTokenizer

model_id = "saugataroy/ielts-task2-modernbert-base-moe"
revision = "main"  # Pin the repository commit SHA for reproducible inference.
tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision)
model = AutoModel.from_pretrained(
    model_id, revision=revision, trust_remote_code=True,
).eval().to("cpu")  # Use "cuda" for NVIDIA or "mps" for a supported Apple GPU.
result = model.score(
    tokenizer,
    questions=["Should cities make public transport free? Discuss both views."],
    essays=[Path("essay.txt").read_text(encoding="utf-8")],
)[0]
print(result)

Review and pin modeling_ielts.py, which trust_remote_code=True loads. Batch inference accepts lists of questions and essays. A generic text-classification pipeline does not create the required essay-token mask; use score or reproduce the paired preprocessing. Low-level forward takes input_ids, attention_mask, and essay_mask, and returns logits shaped [batch, 4] containing regression estimates.

Architecture

Question + essay → paired tokenizer → 22-layer ModernBERT-base
  → mean pooling over contextualized essay tokens
  → four shared experts + four criterion-specific softmax gates
  → criterion-specific linear outputs + fixed 4.5 offset
  → TR, CC, LR, GRA

The encoder uses hidden size 768 and 12 attention heads per layer. Each expert is Linear(768,256) → GELU → Dropout(0.1) → Linear(256,256) → GELU. Each criterion has a Linear(768,4) softmax gate and a Linear(256,1) output. All four experts execute for every essay. This is a dense soft mixture, not sparse top-k MoE routing. The output scores themselves use no softmax or sigmoid. Dropout is disabled in evaluation.

Training data and procedure

The parent encoder was first supervised on ELLIPSE (3,324 training and 587 validation responses, six human-rated 1–5 proficiency traits). A fresh IELTS head and the encoder were then trained on chillies/ielts-writing-task2-essays (7,069 training, 529 validation, 450 test; four Writing9 criteria). ELLIPSE traits were not mapped to IELTS bands. Original splits group related prompts and lexical near-duplicate essays; they cannot guarantee semantic topic isolation.

This continuation initializes all existing encoder and MoE head weights from the pinned parent, without resetting the head. Each epoch uses 976 synthetic training responses plus 1,952 fixed original-training replay responses sampled with seed 42. Replay examples come only from the original training split. Synthetic validation, diagnostic, and quarantine rows are excluded from training. Only question and essay text enter the model; rationales, flags, metadata, and scores are never input text.

Training uses masked mean squared error, AdamW, encoder/head learning rates 5e-6/5e-5, weight decay 0.01 on matrix parameters, gradient clipping at 1.0, 10% linear warm-up, and an effective batch size of 16 (microbatch 2 × accumulation 8). All encoder and head parameters are updated. The run uses seed 42, gradient checkpointing, and CUDA mixed precision. Full environment details are in environment.json.

The run completed 5 epochs within a five-epoch cap; the released checkpoint is continuation epoch 5. Selection minimizes the equal-weight mean of original and synthetic validation macro MAE, subject to an original-validation guard of at most +0.10 macro MAE and +0.20 MAE for any criterion relative to the parent. Early stopping uses two epochs without an eligible gain of at least 0.01. The parent remains eligible as epoch zero. Final diagnostic and test labels are not used for epoch selection. See training_config.json.

Training source and reproducibility

Training and evaluation code is published in ielts-writing-scoring, pinned for this checkpoint at commit ca8b66c.

• Training implementation
• Experiment configuration
• Recorded Kaggle notebook
• Setup, training commands, and reproduction prerequisites
• Source hashes and verification instructions

The notebook's embedded source archive matches the completed run's recorded SHA-256: 46fab05e09f4d4e46537f39c99c30c2ba8144d0f8ef11270c0ee537df76d21d9. The repository includes a script that verifies the current source against that archive without executing notebook code. Exact reruns require the recorded private prepared data; the original transfer stage also requires the intermediate ELLIPSE checkpoint. These are not redistributed by the code repository. Rebuilding data can change Parquet bytes, and the frozen training scripts enforce the recorded checksums. See the linked training guide for these limitations. Model weights and evaluation artifacts remain in this Hugging Face repository.

Evaluation

All errors below use raw, unrounded predictions in band units. Lower MAE and RMSE are better; bias is prediction minus reference. Macro MAE averages the four criterion MAEs. The parent and updated model were evaluated on identical rows with identical preprocessing.

Original validation

529 responses; 23 prompt groups; Writing9 labels. Used for selection.

Criterion	Parent MAE	Updated MAE	Updated RMSE	Updated bias
TR	1.1929	1.2755	1.6182	-0.0267
CC	0.9630	1.0877	1.4091	+0.0871
LR	0.9198	0.9344	1.3210	-0.0593
GRA	1.1115	1.1266	1.5818	-0.0465
Macro	1.0468	1.1061	1.4825	—

Change in macro MAE (updated − parent): +0.0593; paired prompt-group bootstrap 95% interval [+0.0401, +0.0830] (23 groups; 5,000 resamples).

Synthetic validation

98 responses; LLM labels. Used for selection.

Criterion	Parent MAE	Updated MAE	Updated RMSE	Updated bias
TR	2.7470	1.1642	1.5504	-0.4228
CC	2.7242	0.9284	1.4180	-0.2907
LR	1.2776	0.8404	1.2000	-0.0537
GRA	0.9885	0.5467	0.8787	+0.0113
Macro	1.9343	0.8699	1.2618	—

Change in macro MAE (updated − parent): -1.0644; paired prompt-group bootstrap 95% interval [-1.1893, -0.8877] (9 groups; 5,000 resamples).

Original comparison holdout

450 responses; 23 prompt groups; Writing9 labels. Previously evaluated for the parent model.

Criterion	Parent MAE	Updated MAE	Updated RMSE	Updated bias
TR	1.2078	1.3055	1.6596	+0.0979
CC	1.0243	1.1801	1.5027	+0.1900
LR	1.0365	1.0265	1.4375	+0.1223
GRA	1.2469	1.2750	1.7393	+0.0562
Macro	1.1289	1.1968	1.5848	—

Change in macro MAE (updated − parent): +0.0679; paired prompt-group bootstrap 95% interval [+0.0372, +0.0952] (23 groups; 5,000 resamples).

Synthetic diagnostic holdout

120 responses; LLM labels. Excluded from training and epoch selection.

Criterion	Parent MAE	Updated MAE	Updated RMSE	Updated bias
TR	2.6221	1.0776	1.5663	+0.0096
CC	2.7734	1.0590	1.7171	+0.0092
LR	1.2891	0.7961	1.2153	+0.1904
GRA	0.8307	0.4359	0.6776	+0.0883
Macro	1.8788	0.8422	1.2941	—

Change in macro MAE (updated − parent): -1.0367; paired prompt-group bootstrap 95% interval [-1.1806, -0.9118] (11 groups; 5,000 resamples).

Intervals resample prompt groups; they do not capture variability across training seeds or the uncertainty of the labels. Validation intervals are exploratory because those splits selected the checkpoint. The original comparison holdout has already been reported for the parent, so it is not a newly unseen evaluation set.

Public release requires at least 0.15 lower synthetic-diagnostic macro MAE, a paired group interval entirely below zero, and no original-validation or original-test regression beyond +0.10 macro MAE or +0.20 for any criterion. These thresholds were fixed before training. Passing them does not establish reliable examiner grading. Machine-readable results are in evaluation.json.

Technical validation and performance

The exported Transformers package was strictly loaded offline, compared tensor by tensor with the selected training checkpoint, and checked for prediction agreement, truncation reporting, and blank-input rejection. See inference_verification.json.

Production throughput, p95 HTTP latency, and concurrent-serving memory have not been benchmarked for this updated checkpoint. Earlier checkpoint timings should not be presented as measurements of this release. Architecture-level MoE naming does not imply an inference-speed advantage over the plain head.

Licensing and provenance

See LICENSE.md for the checkpoint's licensing-status notice. The base encoder is Apache 2.0; the ELLIPSE corpus declares CC BY-NC-SA 4.0; the IELTS dataset uploader declares CC BY 4.0; the synthetic dataset is CC BY 4.0. These dataset terms do not automatically determine the legal status of trained weights. Public download availability does not establish unrestricted commercial-use rights.

The model is not affiliated with or endorsed by IELTS, Cambridge, British Council, or IDP. Pin the repository commit and report the model and dataset revisions when citing results. provenance.json records the parent, selected epoch, hashes, and source snapshot. Report issues through this repository's Community tab.