MELD: Multi-Task Equilibrated Learning Detector

A 396M-parameter binary AI-vs-human text detector.

The released weights are an SWA average over the top-10 best-AUROC snapshots during training. The architecture is a shared encoder (jhu-clsp/ettin-encoder-400m) with masked-mean pooling and four classification heads. Only the main 2-way head is used at inference; the generator / attack / domain heads are training-time auxiliaries kept in the state dict for completeness.

Files

• model.safetensors — fp32 weights (~1.6 GB).
• meld_config.json — backbone id, head sizes, training hyperparameters, label vocabularies for the auxiliary heads.
• config.json — HuggingFace marker file.
• tokenizer.json, tokenizer_config.json, special_tokens_map.json — copied from the upstream backbone for fully-offline loading.

Inference

The full inference path is below. No external code is required: the snippet below, plus pip install torch transformers safetensors, is sufficient to score a document.

import json
from pathlib import Path

import torch
import torch.nn as nn
from safetensors.torch import load_file
from transformers import AutoModel, AutoTokenizer


class MELDDetector(nn.Module):
    """Shared encoder + four heads. Only `head_main` is used at inference;
    the aux heads remain in the module so the released state_dict loads
    cleanly without missing/unexpected keys."""

    def __init__(self, backbone, n_generators, n_attacks, n_domains,
                 num_labels=2, dropout=0.1):
        super().__init__()
        self.backbone = AutoModel.from_pretrained(
            backbone, attn_implementation="sdpa"
        )
        if hasattr(self.backbone.config, "reference_compile"):
            self.backbone.config.reference_compile = False
        H = self.backbone.config.hidden_size
        self.dropout = nn.Dropout(dropout)
        self.head_main = nn.Sequential(
            nn.Linear(H, H), nn.GELU(), nn.Dropout(dropout),
            nn.Linear(H, num_labels),
        )
        self.head_gen = nn.Linear(H, n_generators)
        self.head_att = nn.Linear(H, n_attacks)
        self.head_dom = nn.Linear(H, n_domains)
        self.log_var_main = nn.Parameter(torch.zeros(()))
        self.log_var_gen = nn.Parameter(torch.zeros(()))
        self.log_var_att = nn.Parameter(torch.zeros(()))
        self.log_var_dom = nn.Parameter(torch.zeros(()))

    def forward(self, input_ids, attention_mask):
        out = self.backbone(input_ids=input_ids,
                            attention_mask=attention_mask).last_hidden_state
        mask = attention_mask.unsqueeze(-1).to(out.dtype)
        pooled = (out * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1.0)
        pooled = self.dropout(pooled)
        return self.head_main(pooled).float()  # (B, 2): [logit_human, logit_ai]


def load_meld(model_dir, device="cpu"):
    cfg = json.loads(Path(f"{model_dir}/meld_config.json").read_text())
    model = MELDDetector(
        backbone=cfg["backbone"],
        n_generators=cfg["n_generators"],
        n_attacks=cfg["n_attacks"],
        n_domains=cfg["n_domains"],
        num_labels=cfg.get("num_labels", 2),
        dropout=cfg.get("dropout", 0.1),
    ).to(device)
    state = load_file(f"{model_dir}/model.safetensors")
    model.load_state_dict(state, strict=True)
    return model.eval(), cfg


# --- usage ---
model_dir = "path/to/this/repo"   # local folder you downloaded
device = "cuda" if torch.cuda.is_available() else "cpu"

model, cfg = load_meld(model_dir, device=device)
tok = AutoTokenizer.from_pretrained(model_dir)   # tokenizer ships in this repo

text = "Your document text goes here."
enc = tok(text, return_tensors="pt", truncation=True,
          max_length=cfg["max_length"]).to(device)
with torch.no_grad():
    meld_score = torch.softmax(model(enc["input_ids"], enc["attention_mask"]),
                               dim=-1)[0, 1].item()
print(f"MELD score (uncalibrated) = {meld_score:.4f}")
# NOTE: a relative ranking score, not a calibrated P(text is AI).
# Compare against a per-pool calibrated threshold — see
# "Intended use and limitations" below.

For documents longer than cfg["max_length"] (2048 tokens), the paper's evaluation protocol uses overlapping 2048-token chunks with mean aggregation of meld_score. A simple stride-512 sliding window over tok(text).input_ids reproduces it.

Intended use and limitations

Research artifact for English document-level AI-text detection; not a forensic tool, and multilingual detection is out of scope.

The output is a relative score optimized for ranking (AUROC / TPR@FPR), not a calibrated per-document probability. Deployment therefore requires threshold calibration on a population matched to the target domain (scripts/calibrate_thresholds.py); a single fixed threshold should not be assumed to transfer across domains, and out-of-distribution or heavily formal text may need domain-specific calibration data, as is standard for detectors of this type.

License

MIT for the model weights, inheriting the upstream backbone license (jhu-clsp/ettin-encoder-400m).