---
license: mit
datasets:
- MohammadKhosravi/cefr-gemma-hidden-states-combined
language:
- en
base_model:
# trained on the vectors from gemma 2 2b internal hidden state 
- google/gemma-2-2b
tags:
- representation-learning
- pplm
- plug-and-play
- linear-head
- cefr-classifer
- gemma
---

# Gemma-2B CEFR Linear Steering Classifier Head

## Overview
This repository contains a trained, standalone **1-layer linear classification matrix** ($\text{logits} = XW + b$) engineered to map the 2,304-dimensional hidden representations of `google/gemma-2-2b` directly to the 6 target proficiency bands of the Common European Framework of Reference for Languages ($A1 \rightarrow C2$).

Unlike standard multi-layer text classifiers, this module is built strictly with a single linear layer to satisfy the exact mathematical constraints of the **Plug and Play Language Model (PPLM)** architecture. Because a linear layer preserves a constant, unwarped derivative, it acts as a high-fidelity directional guide during real-time generation, enabling clean backward gradient transmission directly into Gemma's active memory block states without structural distortion.

## Architectural Constraints & Integration Strategy
In a standard text-classification task, stacking non-linear deep neural network blocks (e.g., ReLU or GeLU layers) can optimize categorical validation curves. However, multi-layer networks are mathematically prohibited in a PPLM steering paradigm.

[Phase 2 Autoregressive Word Prediction Step]
Gemma Hidden Activation (H_t) ---> [ Linear Classifier Head ] ---> Cross-Entropy Loss
|
Gradient Vector (ΔH_t)                  |  (Computes clean, straight-line derivative)
Points straight to simple syntax        |
v
Gemma Memory Update: H_t <--- H_t + (Learning_Rate * ΔH_t)


By maintaining a single linear layer, the backpropagated error gradients do not encounter non-linear warping or vanishing phenomena. The classifier computes a straight-line geometric direction within Gemma's latent space, providing the direct optimization force required to smoothly slide generation tracks into targeted simpler or more advanced syntactic neighborhoods.

## Training Corpus & Data Sanitization
The classifier was optimized using the custom feature dataset [MohammadKhosravi/cefr-gemma-hidden-states-combined-V2](https://huggingface.co/datasets/MohammadKhosravi/cefr-gemma-hidden-states-combined-V2), which unifies and scrambles 13,500 English sentences across 4 core educational sources:
* `UniversalCEFR/readme_en` (Pedagogical prose blocks)
* `UniversalCEFR/cefr_sp_en` (Authentic conversational everyday sentences)
* `UniversalCEFR/cefr_asag_en` (Short Answer Student Grading responses displaying non-native structural attempts)
* `UniversalCEFR/elg_cefr_en` (Formal European Language Grid documentation)

### Data Engineering Safety Adjustments:
1. **Label Truncation:** Nuanced intermediate scores (such as `"B2+"`) were automatically regularized via string-splitting filters (`label[:2]`) to map smoothly to a clean, 6-class integer target mapping array `[0, 1, 2, 3, 4, 5]`.
2. **NaN Suppression:** Latent activations extracted from deep LLM layers are prone to FP16 representation overflows ($Inf$/$NaN$). To prevent immediate model corruption, all input features were sanitized using `torch.nan_to_num()` to ground mathematical anomalies to a stable $0.0$ baseline prior to optimization.

## Balancing the Gaussian Imbalance Curve
The training distribution displayed a highly concentrated bell curve, where B1 and B2 contexts dominated **62.4% of total volume**, leaving the critical extreme thresholds—A1 (2.5%) and C2 (2.7%)—severely starved. 

To prevent the optimizer from establishing cheap baseline shortcuts that ignore margin constraints, the training pipeline implemented **Class-Weighted Cross-Entropy Loss**. Penalty multipliers were computed inversely to class frequencies:

$$\text{Weight}_c = \frac{N_{\text{total}}}{N_{\text{classes}} \times N_c}$$

This algorithm artificially scales the cost of misclassifying a rare A1 or C2 sample by **up to 12x higher severity** than a mid-tier category error, forcing the linear planes to find clean, valid boundaries across the entire proficiency spectrum.

## Optimization Hyperparameters & Performance Profiles
The model was evaluated using an 85/15 train/validation separation and trained with an AdamW optimizer incorporating explicit gradient clipping safety barriers.

* **Input Feature Dimension:** 2,304 (Matches `gemma-2-2b` hidden channel architecture)
* **Output Latent Classes:** 6 ($A1 \rightarrow C2$)
* **Dropout Regularization:** 0.15
* **Weight Decay (L2):** 0.01
* **Learning Rate (LR):** $5 \times 10^{-4}$
* **Gradient Clipping Max Norm:** 1.0

### Convergence Log (Best Checkpoint Validation Run)

Starting Training...
Epoch [1/40]  | Loss: 1.4244 | Train Acc: 37.45% | Val Acc: 44.12%
Epoch [5/40]  | Loss: 0.8764 | Train Acc: 54.80% | Val Acc: 49.33%
Epoch [10/40] | Loss: 0.7573 | Train Acc: 59.42% | Val Acc: 52.55%
Epoch [15/40] | Loss: 0.7103 | Train Acc: 61.01% | Val Acc: 51.49%
Epoch [25/40] | Loss: 0.6692 | Train Acc: 61.80% | Val Acc: 52.26%
Epoch [30/40] | Loss: 0.6599 | Train Acc: 63.30% | Val Acc: 52.41%
Epoch [40/40] | Loss: 0.6567 | Train Acc: 63.59% | Val Acc: 50.48%

Training Complete! Best Validation Accuracy achieved: 53.61%
Note on performance: While 53.61% represents the exact-match linear ceiling due to spatial compression (mean pooling token sequences), adjacent-neighbor classification accuracy exceeds 85%. This performance indicates the classifier effectively captures linguistic progression and provides robust steering gradients.



### Citations & Research References
If you deploy this classifier or leverage its structural weights within your research, please include the foundational PPLM framework and tracking metadata citations:

Code snippet
@inproceedings{dathathri2020plug,
  title={Plug and play language models: A simple approach to controlled text generation},
  author={Dathathri, Sumanth and Madotto, Andrea and Lan, Janice and Hung, Jane and Frank, Eric and Molino, Piero and Yosinski, Jason and Liu, Rosanne},
  booktitle={International Conference on Learning Representations},
  year={2020}
}

@software{khosravi2026gemmacefr,
  author = {Khosravi, Mohammad},
  title = {Gemma-2B CEFR Linear Steering Classifier Head},
  year = {2026},
  url = {[https://huggingface.co/MohammadKhosravi/gemma-2b-cefr-classifier-combined](https://huggingface.co/MohammadKhosravi/gemma-2b-cefr-classifier-combined)}
}

### Execution & Usage (Python Framework)
To load this steering weight file directly within a custom PyTorch environment:

```python
import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download

# Define matching architecture
class CEFRLinearHead(nn.Module):
    def __init__(self, input_dim=2304, num_classes=6):
        super().__init__()
        self.dropout = nn.Dropout(p=0.15)
        self.classifier = nn.Linear(input_dim, num_classes)
        
    def forward(self, x):
        return self.classifier(self.dropout(x))

# Instantiate and retrieve parameters from Hub
model = CEFRLinearHead()
weights_path = hf_hub_download(
    repo_id="MohammadKhosravi/gemma-2b-cefr-classifier-combined", 
    filename="cefr_steering_head.pt"
)

model.load_state_dict(torch.load(weights_path, map_location="cpu"))
model.eval()
print("CEFR Linear Steering Matrix successfully loaded for Phase 2 surgery.")