MohammadKhosravi/llama3.1-8b-pure-pmt-cefr-gating-6k

Pure PrefixMemory-Tuning (PMT) + CEFR Gating

This model implements a highly faithful replication of the uncompressed PrefixMemory-Tuning architecture (arXiv:2506.13674), modified for explicit CEFR-aligned text generation. By replacing low-rank bottlenecks with full-dimensional layer-wise memory matrices and introducing a feature-wise CEFR gating mechanism, it achieves state-of-the-art parameter-efficient controllability. The model steers linguistic complexity architecturally, meaning the desired CEFR proficiency level is enforced directly through learned parameters rather than textual prompt instructions.

Base Model

• Foundation Model: Meta-Llama-3.1-8B-Instruct
• Parameters: 8 Billion
• Backbone: Decoder-only Transformer
• Transformer Layers: 32
• Hidden Dimension: 4096
• Frozen during training: Yes. Only the PrefixMemory controller is trainable.

---

Architecture

This iteration eliminates the low-rank bottleneck projections and instead uses a fully uncompressed memory parameterization.

Specifically, the model removes:

$$W_{\mathrm{d}\mathrm{o}\mathrm{w}\mathrm{n}}\text{and}{W}_{\mathrm{u}\mathrm{p}}W\_\{\backslash mathrm\{down\}\}\; \backslash quad\; \backslash text\{and\}\; \backslash quad\; W\_\{\backslash mathrm\{up\}\}$$

and replaces them with a direct memory matrix formulation.

The resulting controller contains approximately

$$536.9\text{ million}536.9\; \backslash text\{\; million\}$$

trainable parameters.

A trainable CEFR embedding

$$e_{\mathrm{c}\mathrm{e}\mathrm{f}\mathrm{r}}\in {\mathbb{R}}^{4096}e\_\{\backslash mathrm\{cefr\}\}\; \backslash in\; \backslash mathbb\{R\}^\{4096\}$$

is used to perform feature-wise gating of the activated hidden states. The embedding is applied through element-wise modulation of the ELU feature map before querying the layer-specific uncompressed memory matrix.

This design removes the low-rank information bottleneck and allows CEFR control signals to directly influence all 4096 hidden dimensions at every transformer layer.

Figure 1. Overall architecture of the Pure PrefixMemory-Tuning (PMT) + CEFR Gating

Controller Pipeline

For each transformer layer $l \in [0, 31]$:

1. Kernel Feature Mapping

Hidden states $X$ are passed through an Exponential Linear Unit (ELU):

$$\varphi (X)=\mathrm{E}\mathrm{L}\mathrm{U}(X)\backslash phi(X)\; =\; \backslash mathrm\{ELU\}(X)$$

2. Feature-Wise CEFR Gating

The target CEFR embedding vector

$$e_{\mathrm{c}\mathrm{e}\mathrm{f}\mathrm{r}}\in {\mathbb{R}}^{4096}e\_\{\backslash mathrm\{cefr\}\}\; \backslash in\; \backslash mathbb\{R\}^\{4096\}$$

is broadcasted and multiplied element-wise with the activated hidden states:

$$\tilde{X}=\varphi (X)\odot e_{\mathrm{c}\mathrm{e}\mathrm{f}\mathrm{r}}\backslash widetilde\{X\}\; =\; \backslash phi(X)\; \backslash odot\; e\_\{\backslash mathrm\{cefr\}\}$$

3. Pure Memory Retrieval

The modulated features are multiplied by the layer-specific uncompressed memory matrix:

$$M^{(l)}\in {\mathbb{R}}^{4096\times 4096}M^\{(l)\}\; \backslash in\; \backslash mathbb\{R\}^\{4096\; \backslash times\; 4096\}$$

yielding

$${\mathrm{B}\mathrm{i}\mathrm{a}\mathrm{s}}_{\mathrm{m}\mathrm{e}\mathrm{m}}=\tilde{X}{M}^{(l)}\backslash mathrm\{Bias\}\_\{\backslash mathrm\{mem\}\}\; =\; \backslash widetilde\{X\}\; M^\{(l)\}$$

4. Attention Injection

The memory bias is scaled by a learnable parameter $\alpha$ and injected directly into the frozen Llama attention output:

$${\mathrm{O}\mathrm{u}\mathrm{t}\mathrm{p}\mathrm{u}\mathrm{t}}_{\mathrm{f}\mathrm{i}\mathrm{n}\mathrm{a}\mathrm{l}}={\mathrm{O}\mathrm{u}\mathrm{t}\mathrm{p}\mathrm{u}\mathrm{t}}_{\mathrm{a}\mathrm{t}\mathrm{t}\mathrm{n}}+\alpha {\mathrm{B}\mathrm{i}\mathrm{a}\mathrm{s}}_{\mathrm{m}\mathrm{e}\mathrm{m}}\backslash mathrm\{Output\}\_\{\backslash mathrm\{final\}\}\; =\; \backslash mathrm\{Output\}\_\{\backslash mathrm\{attn\}\}\; +\; \backslash alpha\; \backslash ,\; \backslash mathrm\{Bias\}\_\{\backslash mathrm\{mem\}\}$$

---

Training Objective

The controller was trained using the standard cross-entropy causal language modeling objective via teacher forcing.

The entire Meta-Llama-3.1-8B-Instruct backbone remained strictly frozen throughout training.

Gradients flowed exclusively into:

• CEFR embedding parameters
• Layer-specific memory matrices $M^{(l)}$

The objective is to maximize the likelihood of human-written CEFR-aligned reference texts:

where $c$ denotes the target CEFR level.

---

Training Dataset

Training uses a balanced subset derived from the EFCAMDAT learner corpus.

Original Dataset

EFCAMDAT (Cambridge English)

Due to the EFCAMDAT data usage policy, the processed training subset cannot be redistributed.

Dataset Columns

Column	Description
topic_title	Writing prompt
clean_text	Human learner response
cefr	Target CEFR level

---

Distribution

Balanced across all six CEFR proficiency levels.

CEFR	Samples
A1	1,000
A2	1,000
B1	1,000
B2	1,000
C1	1,000
C2	1,000

Total Samples: 6,000

---

Training Performance

Best checkpoint Saved:

Epoch 1

[
  {
    "epoch": 1,
    "train_loss": 2.4011,
    "val_loss": 2.0971,
    "val_ppl": 8.14
  },
  {
    "epoch": 2,
    "train_loss": 1.6818,
    "val_loss": 2.1277,
    "val_ppl": 8.40
  },
  {
    "epoch": 3,
    "train_loss": 1.0301,
    "val_loss": 2.4258,
    "val_ppl": 11.31
  },
  {
    "epoch": 4,
    "train_loss": 0.5113,
    "val_loss": 2.7793,
    "val_ppl": 16.11
  },
  {
    "epoch": 5,
    "train_loss": 0.2538,
    "val_loss": 3.1331,
    "val_ppl": 22.94
  }
]

Checkpoint Selection

Epoch 1 is used for inference.

Although later epochs achieve lower training loss, validation perplexity steadily increases, indicating overfitting of the uncompressed 536 million parameter PrefixMemory controller. The Epoch 1 checkpoint provides the best balance between controllability and generalization.

In-Domain Evaluation

Benchmark

A balanced benchmark constructed from completely held-out EFCAMDAT prompts.

Benchmark Statistics

• Total prompts: 702
• Evaluation style: Prompt-only (Blind)
• Target CEFR instructions in prompt: None
• Distribution: 117 prompts per CEFR level

---

Improvement over Base Llama (In-Domain)

Compared to the frozen Meta-Llama-3.1-8B-Instruct baseline:

• +44.16% Strict Accuracy

  • Improved from 24.64% to 68.80%

• +18.15% Adjacent Accuracy

  • Improved from 63.19% to 81.34%

• Halved Prediction Error

  • Mean Absolute Error (MAE) reduced from 1.3275 to 0.6368

• Natural Readability Maintained

  • Average Flesch Reading Ease: 82.44

• Pure Architectural Control

  • No CEFR instructions were included in the text prompt.
  • The model followed the internal CEFR control signal exclusively, demonstrating genuine parameter-level controllability rather than prompt-following behavior.

---

PMT Evaluation Results

============================================================
📊 FINAL COMPILED MACRO-STATISTICS (PURE PMT + CEFR GATING)
============================================================

Total Processed       : 702
Strict Accuracy       : 68.80%
Adjacent Accuracy     : 81.34%
Mean Abs Error (MAE)  : 0.6368
Avg MDD Score         : 1.86
Avg Reading Ease      : 82.44
Avg Sentence Drift    : 2.05 levels

============================================================
📈 PERFORMANCE BREAKDOWN BY CEFR TARGET LEVEL
============================================================

             strict_accuracy  avg_mdd  avg_flesch  avg_drift
target_cefr
-------------------------------------------------------------
A1                     85.47     1.42       81.12       0.56
A2                     86.32     2.16       91.05       1.03
B1                     76.92     1.98       81.09       2.04
B2                     68.38     1.78       82.50       2.95
C1                     51.28     1.82       79.00       2.68
C2                     44.44     2.00       79.90       3.02

============================================================
📝 DETAILED CLASSIFICATION REPORT
============================================================

              precision    recall  f1-score   support

A1               0.71       0.85      0.78       117
A2               0.71       0.86      0.78       117
B1               0.60       0.77      0.68       117
B2               0.55       0.68      0.61       117
C1               0.87       0.51      0.65       117
C2               0.93       0.44      0.60       117

-------------------------------------------------------------

Accuracy                                 0.69       702
Macro Avg         0.73       0.69      0.68       702
Weighted Avg      0.73       0.69      0.68       702

Confusion Matrix

Figure 2. Confusion Matrix of PrefixMemory Tuning on In-Domain Evaluation Dataset (Inference).

---

Baseline LLM Evaluation Results (Reference)

============================================================
📊 FINAL COMPILED MACRO-STATISTICS (BASE LLM)
============================================================

Total Processed       : 690
Strict Accuracy       : 24.64%
Adjacent Accuracy     : 63.19%
Mean Abs Error (MAE)  : 1.3275
Avg MDD Score         : 2.00
Avg Reading Ease      : 49.47
Avg Sentence Drift    : 3.04 levels

============================================================
📈 PERFORMANCE BREAKDOWN BY CEFR TARGET LEVEL
============================================================

             strict_accuracy  avg_mdd  avg_flesch  avg_drift
target_cefr
-------------------------------------------------------------
A1                     16.24     1.73       82.56       3.25
A2                     23.08     1.64       85.44       3.27
B1                     17.39     2.03       57.80       3.11
B2                     41.74     2.15       34.18       2.97
C1                     33.63     2.24       18.95       2.94
C2                     15.93     2.25       15.54       2.69

Confusion Matrix

Figure 3. Confusion Matrix of Base LLM on In-Domain Evaluation Dataset (Inference).

---

Out-of-Domain Evaluation

Benchmark

Dataset:

MohammadKhosravi/ielts-cefr-benchmark-prompts

A manually curated benchmark containing IELTS-style prompts designed to evaluate out-of-domain CEFR controllability.

Dataset Characteristics

• Balanced distribution
• 50 prompts per CEFR level
• 300 prompts total
• Completely unseen prompt domain

---

Improvement over Base Llama (OOD)

Compared to the frozen baseline:

• +40.00% Strict Accuracy

  • Improved from 18.33% to 58.33%

• +21.34% Adjacent Accuracy

  • Improved from 54.33% to 75.67%

• Robust Generalization

  • Full-rank memory matrices ($4096 \times 4096$) successfully transferred CEFR control to unseen IELTS-style prompts.

• Substantial Error Reduction

  • MAE reduced from 1.5733 to 0.7700

---

OOD Evaluation Results

============================================================
📊 FINAL COMPILED MACRO-STATISTICS
(PURE PMT + EMBEDDING GATING - IELTS OOD)
============================================================

Total Processed       : 300
Strict Accuracy       : 58.33%
Adjacent Accuracy     : 75.67%
Mean Abs Error (MAE)  : 0.7700
Avg MDD Score         : 1.88
Avg Reading Ease      : 78.13
Avg Sentence Drift    : 2.08 levels

============================================================
📈 PERFORMANCE BREAKDOWN BY CEFR TARGET LEVEL
============================================================

             strict_accuracy  avg_mdd  avg_flesch  avg_drift
target_cefr
-------------------------------------------------------------
A1                      60.0     1.64       80.38       1.08
A2                      86.0     1.79       93.90       0.74
B1                      70.0     2.02       71.64       2.02
B2                      48.0     1.88       72.88       2.72
C1                      64.0     1.93       74.00       2.92
C2                      22.0     2.03       75.99       2.98

OOD Confusion Matrix

---

Intended Use

This model is intended for research on:

• Controllable text generation
• CEFR-level writing generation
• Educational NLP
• Readability control
• Parameter-Efficient Fine-Tuning (PEFT)
• PrefixMemory-Tuning

---

Citation

If you use this model, please cite:

• PrefixMemory-Tuning (arXiv:2506.13674)
• Meta-Llama-3.1-8B-Instruct
• EFCAMDAT Corpus
• This repository