Llama-SARM-4B

SARM: Interpretable Reward Model via Sparse Autoencoder

This repository contains the model weights of the AAAI 2026 Oral Paper "Interpretable Reward Model via Sparse Autoencoder".

🔥 News

• [2025/11/8] Our paper has been accepted as an oral presentation at AAAI 2026. 🎉
• [2025/12/11] Llama-SARM-4B is ranked 18th on the Reward Bench 2 leaderboard, above GPT-4.1, Skywork-Reward-Llama-3.1-8B, and Claude-Sonnet-4!🎉

🔗 Links

• Authors

  Shuyi Zhang*, Wei Shi*, Sihang Li*, Jiayi Liao, Tao Liang, Hengxing Cai, Xiang Wang†

• Paper: Interpretable Reward Model via Sparse Autoencoder

• Code Repository: https://github.com/schrieffer-z/sarm

• Demo: Try SARM Demo in Huggingface Space

📊 Evaluation

Llama-SARM-4B shows competitive performance, even with a much smaller parameter size.

Reward Bench 2

SARM Inference Demo

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer


def get_reward_score(model, prompt, response) -> float:
    """
    Receives a prompt and a response, and returns the reward score calculated by the SARM model.
    """
    messages = [{"role": "user", "content": prompt}, {"role": "assistant", "content": response}]
    input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device)

    with torch.no_grad():
        score = model(input_ids).logits.item()

    return round(score, 4)


device = "cuda"
path = "Schrieffer/Llama-SARM-4B"

tokenizer = AutoTokenizer.from_pretrained(path)
model = AutoModelForSequenceClassification.from_pretrained(
    path, 
    device_map=device, 
    trust_remote_code=True, 
    torch_dtype=torch.bfloat16
)

examples=[
    ["What is the capital of France?", "The capital of France is Paris."],
    ["What is the capital of France?", "Berlin is a large city in Germany."],
    ["Write a short poem about the moon.", "Silver orb in velvet night, / Casting shadows, soft and light. / Silent watcher, distant, bright, / Guiding dreams till morning's light."],
    ["Write a short poem about the moon.", "The moon is a rock."]
]

for example in examples:
    print("example".center(80,'='))
    print("Question:
"+example[0])
    print("Answer:
"+example[1])
    print("Score:", get_reward_score(model, example[0],example[1]))

📧 Contact

If you have any questions, please feel free to reach us at [email protected].

📚 Citation

If you find our work useful, please cite it as follows.

@article{zhang2025interpretable,
  title={Interpretable Reward Model via Sparse Autoencoder},
  author={Zhang, Shuyi and Shi, Wei and Li, Sihang and Liao, Jiayi and Liang, Tao and Cai, Hengxing and Wang, Xiang},
  journal={arXiv preprint arXiv:2508.08746},
  year={2025}
}