ai-detection-demo-gemma_4_e4b

ai-detection-demo — Gemma 4 E4B linear probe

A tiny (~1.6 M param) classification head trained on top of frozen google/gemma-4-E4B-it last-token logits, for AI-vs-human text classification.

Same task and label convention as the EditLens paper, except the backbone is never adapted — only this tiny probe is trained.

Architecture

raw_logits[262144] → LayerNorm → Dropout(0.3) → Linear → cls_logits[4]

Total params: ~1.57 M.

Files

Usage

import json
import torch
from safetensors.torch import load_file
from huggingface_hub import snapshot_download

# See ai-detection-demo repo for the LinearDropoutHead class definition.
from gemma_4_e4b.head import LinearDropoutHead, score_from_logits

local = snapshot_download("DarrenJiaImbue/ai-detection-demo-gemma_4_e4b")
with open(f"{local}/config.json") as f:
    cfg = json.load(f)

head = LinearDropoutHead(
    vocab_size=cfg["vocab_size"], n_classes=cfg["n_classes"], dropout=cfg["dropout"]
)
head.load_state_dict(load_file(f"{local}/head.safetensors"))
head.eval()

# raw_logits: (batch, 262144) fp32 tensor of Gemma's last-token vocabulary logits.
with torch.no_grad():
    logits = head(raw_logits)
    bucket = logits.argmax(dim=-1)          # 0..3
    ai_score = score_from_logits(logits)    # continuous [0, 1]

Training data

Trained on DarrenJiaImbue/ai-detection-demo-gemma-logits — pre-computed int4 Gemma 4 E4B logits over DarrenJiaImbue/ai-detection-demo-dataset.

License

CC BY-NC-SA 4.0. Non-commercial research use only.