You need to agree to share your Hugging Face information to access this model

This repository is publicly accessible, but you have to accept the conditions to access its files and content.

Access is granted by manual approval. Please say briefly who you are and how you plan to use the model.

Write like Kaiming and Saining

LoRA adapters for the pretrained base model google/gemma-4-31B that rewrite a paragraph of an ML paper in the writing style of Kaiming He or Saining Xie.

You give the model a draft paragraph. It returns a rewrite of that paragraph with the same content. The model aims for short direct sentences, clear transitions and concrete claims.

Base model only. We trained these adapters on the pretrained base model google/gemma-4-31B. They do not work with the instruction-tuned model (google/gemma-4-31B-it) or with any other checkpoint. Load the adapters on google/gemma-4-31B only.

This is a personal research project. Kaiming He and Saining Xie did not take part in it, and they do not endorse it.

Versions

The repository holds four adapters. All of them use the same base model, the same prompt format and the same LoRA shape.

Recommended: v6-epoch4/. We trained it with the v5g recipe on a filtered data set (see Training). It changes facts less often than v5g-epoch2 and keeps the paragraph structure more reliably. In a blind side-by-side reading of 24 held-out passages, we preferred v6-epoch4 over v5g-epoch2 in 16 passages and v5g-epoch2 in 7 (1 tie). One weakness: when the input is long and wordy, v6-epoch4 often keeps most of it instead of shortening it. For that case, v5g-epoch2 can give better results.

The older adapters differ in rewrite strength, from light to strong:

Strength	Folder	What it does	When to use it
Light	`v5g-epoch1/`	Makes conservative edits, mostly at the word level. The sentence structure usually stays the same.	You want light polishing and a low risk of changes in meaning.
Medium	`v5g-epoch2/`	Rewrites full sentences and paragraphs. It makes fewer factual errors than `v5f`.	You want to shorten a long, wordy draft.
Strong	`v5f/`	Makes the largest changes and restructures whole paragraphs. It sometimes changes facts.	You want bold alternatives and you will check every claim.

A stronger adapter gives more varied and more fluent rewrites, but it also has a higher risk of changes in meaning.

Training

Base model: google/gemma-4-31B (pretrained base, not the instruction-tuned model).
LoRA: rank 64, alpha 64, on all attention and MLP projections of the language model (q/k/v/o_proj, gate/up/down_proj).
Optimizer: learning rate 2e-4, context length 32k tokens.
Data: papers by Kaiming He and Saining Xie from arXiv. Each training pair maps an AI-written version of a passage to the original text by the author. The passages are sentences, paragraphs, sections and full papers. The data set has approximately 22k pairs.
Epochs: v5f trained for 4 epochs. v5g used the same recipe and we saved it after epoch 1 and epoch 2.
v6-epoch4: the v5g recipe for 4 epochs, on a filtered subset of the pairs. We kept a pair only if its numbers, citations, references and math are identical on both sides, it has the same number of paragraphs on both sides, and the target length is 0.80 to 1.25 times the input length. This kept approximately 11.6k of the 22k pairs. The filter removes pairs that teach the model to add or delete content.

The model was not trained on chat or instruction data. Use the fixed prompt below; free-form instructions do not work.

Prompt format

<style: {person}>
Rewrite the following text.

{your paragraph}

{person} is Kaiming He or Saining Xie. Put the BOS token first. Do not add a chat template. The model writes the rewrite after the blank line and stops with EOS.

Recommended sampling: temperature=0.7, top_p=0.95, top_k=0. Generate 8 to 32 samples per paragraph and select the best ones.

Usage

The adapters are plain safetensors files (lora_state.safetensors) with a small lora_meta.json. They are not in PEFT format. The script below merges an adapter into the base weights and then generates text with transformers. The base model in bf16 needs approximately 62 GB of GPU memory.

import json, re, torch
from huggingface_hub import snapshot_download
from safetensors.torch import load_file
from transformers import AutoModelForCausalLM, AutoTokenizer

REPO = "yipengsun/write-like-kaiming-and-saining"
VERSION = "v6-epoch4"  # or "v5g-epoch2", "v5g-epoch1", "v5f"

lora_dir = snapshot_download(REPO, allow_patterns=[f"{VERSION}/*"]) + f"/{VERSION}"
tok = AutoTokenizer.from_pretrained("google/gemma-4-31B")
model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-4-31B", torch_dtype=torch.bfloat16, device_map="auto")

# Merge the LoRA: W <- W + (alpha / rank) * B @ A
meta = json.load(open(f"{lora_dir}/lora_meta.json"))
sd = load_file(f"{lora_dir}/lora_state.safetensors")
scale = meta["alpha"] / meta["rank"]
modules = dict(model.named_modules())
wrap = re.compile(r"(_fsdp_wrapped_module\.|_checkpoint_wrapped_module\.|_orig_mod\.|_forward_module\.)")
pairs = {}
for key, value in sd.items():
    key = wrap.sub("", key).replace(".default.weight", ".weight")
    side = "A" if key.endswith(".lora_A.weight") else "B"
    name = key.rsplit(".lora_", 1)[0]
    while name not in modules:  # strip outer prefixes until the name matches
        name = name.split(".", 1)[1]
    pairs.setdefault(name, {})[side] = value
assert len(pairs) == meta["n_tensors"] // 2
with torch.no_grad():
    for name, ab in pairs.items():
        w = modules[name].weight
        delta = ab["B"].to(w.device).float() @ ab["A"].to(w.device).float()
        w.copy_((w.float() + scale * delta).to(w.dtype))

draft = "In this paper we propose a novel method that is able to significantly improve the performance of ..."
prompt = f"<style: Kaiming He>\nRewrite the following text.\n\n{draft}\n\n"
ids = [tok.bos_token_id] + tok(prompt, add_special_tokens=False)["input_ids"]
ids = torch.tensor([ids], device=model.device)
out = model.generate(ids, max_new_tokens=2 * ids.shape[1], do_sample=True,
                     temperature=0.7, top_p=0.95, top_k=0, num_return_sequences=4)
for seq in out:
    print(tok.decode(seq[ids.shape[1]:], skip_special_tokens=True), "\n---")

To serve the model with vLLM or another engine, save the merged model with model.save_pretrained(...) and load that folder.

Limitations

Read every output before you use it. The model is a writing aid. It does not check facts.

Invented content. The model can change numbers, change definitions, remove qualifiers or add claims that the input does not support. This happens most often with v5f and with long inputs.
Copying. If the input is already well written, the model often returns it almost unchanged.
Long inputs. When you give a full paper, the model often stops before the conclusion. Some samples collapse to a single short paragraph. Rewrite one section or a few paragraphs at a time for the best results.
LaTeX. The model sometimes breaks math, citations or special characters. Compare the output with the input.
Scope. The training data contains only computer vision and machine learning papers by two authors. Other fields and other styles were not tested.

License

The adapters are released under Apache 2.0, the same license as the base model. Use of the base model follows its own license terms.

Downloads last month: -; Downloads are not tracked for this model. How to track

Model tree for yipengsun/write-like-kaiming-and-saining

Base model

google/gemma-4-31B

Adapter

(4)

this model