jfleg-gec-transformer-ensemble

---

Grammar Error Correction with Copy-Augmented Transformer Ensemble

This repository hosts a progression of custom-built, from-scratch TensorFlow/Keras Transformer models trained specifically for Grammatical Error Correction (GEC) on the JFLEG dataset.

The project demonstrates how to effectively train deep sequence-to-sequence architectures in a severely low-resource environment (approximately 1,200 training pairs) through architectural optimizations, dynamic synthetic data generation, and weighted ensemble methods. The final ensemble model achieves a test accuracy of 72.88% and a corpus GLEU score of 0.5506, representing a massive 3x performance improvement over a standard baseline Transformer.

Model Description

The repository tracks the evolution of the model across three main architectural iterations, proving that architectural efficiency and data pipeline engineering supersede raw parameter count in low-resource NLP:

1. V1 (Baseline): A standard Encoder-Decoder Transformer architecture.
2. V2 (Generative): Incorporates Pre-Layer Normalization (Pre-LN) for gradient stability, GELU activations, and Weight Tying (sharing the embedding matrix across the encoder, decoder, and output projection). Weight tying effectively halves the parameter count (~3.3M parameters), acting as a powerful regularizer against overfitting on small datasets.
3. V3 (Pointer-Generator): Introduces a copy mechanism. Token-level error analysis indicated that approximately 62% of errors occurred on tokens that simply needed to be copied directly from the source sentence. The V3 architecture adds a decoder-to-encoder attention head and a sigmoid gating mechanism. This allows the model to explicitly point to and copy out-of-vocabulary (OOV) and structural tokens from the input sequence, rather than forcing the generative head to guess them.

Training Data & Augmentation Pipeline

To overcome the strict 1,200-sentence constraint of the JFLEG training split, a highly optimized data augmentation pipeline was implemented:

• Multi-Reference Augmentation: The JFLEG dataset contains 4 reference corrections per sentence. All valid references were utilized during training, immediately multiplying the static dataset size by four.
• Identity Pairs: Corrected sentences were fed into the model with themselves as targets to enforce and stabilize the model's copying behavior.
• Dynamic Mined Noise (The Core Engine): Instead of utilizing generic, rule-based noise (which the model easily memorized), error patterns (insertions, deletions, and specific word substitutions) were mined directly from the alignment differences of the actual incorrect-to-correct training pairs. This noise is generated dynamically at each epoch. Consequently, the model never sees the exact same erroneous sentence twice. This simulates a massive dataset (yielding 67,200 unique training pairs per epoch) and prevents catastrophic overfitting.

Experiment Log & Results

All models were evaluated on the exact same unseen test split (151 sentences) using masked categorical accuracy and sparse categorical cross-entropy loss.

Note: The final GLEU score for the Exp 11 Ensemble is 0.5506 (evaluated via batch greedy decoding against 4 references).

Files and Repository Contents

This repository includes the saved weights for the key experimental milestones, allowing researchers to load and test specific phases of the architectural evolution.

• tokenizer_jfleg.json & tokenizer_jfleg.pkl: The fitted vocabulary tokenizer (Vocab size: 12,000).
• best_model_config.json: Configuration hyperparameters, metadata, and ensemble weights.
• transformer_jfleg.weights.h5 to transformer_jfleg_p5.weights.h5: Early iteration baseline weights.
• transformer_jfleg_p6.weights.h5: V2 Architecture trained with dynamic noise.
• transformer_jfleg_p7.weights.h5: V3 Pointer-Generator model.
• transformer_jfleg_p9.weights.h5: V3b Pointer-Generator model (optimized with label smoothing).
• transformer_jfleg_p11a.weights.h5 & transformer_jfleg_p11b.weights.h5: Additional high-capacity pointer models trained with different seeds for the final ensemble.

Intended Use and Limitations

Intended Use: This model is intended for educational and research purposes to demonstrate the viability of training complex sequence-to-sequence models on extremely small datasets using architectural modifications and dynamic data augmentation.

Limitations: The accuracy plateau observed around 73% is primarily an artifact of the dataset mapping rather than model capability constraints. Error decomposition shows that the remaining errors are heavily concentrated on "novel" tokens—words inserted by human annotators that do not exist anywhere in the source sentence. Because the automated test script evaluates exact token matches, the model is inherently penalized for providing valid, alternative grammatical corrections that do not exactly match the specific annotator's arbitrary word choice.

How to Use

Because these are custom-built Keras layers and not standard Hugging Face transformers classes, you must recreate the architecture in your Python environment before loading the weights.

import tensorflow as tf
import pickle
import json

# 1. Load Tokenizer
with open("tokenizer_jfleg.pkl", "rb") as f:
    tokenizer = pickle.load(f)

# 2. Define Custom Classes
# Ensure you have defined the custom classes: 
# PreLNEncoder, PreLNDecoder, TiedOutputLayer, PointerGeneratorHead, TransformerV3b
# (Refer to the training notebook for the exact class definitions).

# 3. Instantiate Model
model_p9 = TransformerV3b(
    num_layers=6, 
    d_model=256, 
    num_heads=8, 
    dff=1024, 
    vocab_size=12000, 
    pe_input=80, 
    pe_target=79,
    dropout_rate=0.2
)

# 4. Build the graph with a dummy pass
dummy_enc_input = tf.constant([[1, 2, 3]])
dummy_dec_input = tf.constant([[1]])
model_p9(dummy_enc_input, dummy_dec_input, training=False)

# 5. Load Weights
model_p9.load_weights("transformer_jfleg_p9.weights.h5")

print("Model successfully loaded and ready for inference!")

Citation & Acknowledgments

If you utilize this training methodology, the dynamic mined noise pipeline, or the pointer-generator modifications, please reference this repository.