Monet: Mixture of Monosemantic Experts for Transformers

Model Summary

Monet introduces a novel approach for improving mechanistic interpretability in large language models (LLMs) using a Sparse Mixture-of-Experts (SMoE) architecture with 262,144 experts. By integrating sparse dictionary learning directly into end-to-end pretraining, Monet tackles the core issue of polysemanticity—where single neurons encode multiple unrelated concepts—while preserving overall model performance.

Resources and Technical Documentation

• GitHub Repository: https://github.com/dmis-lab/Monet
• Paper: https://arxiv.org/abs/2412.04139
• Model Hub: https://huggingface.co/MonetLLM
• Demo: https://huggingface.co/spaces/MonetLLM/monet-vd-1.4B-100BT-hf-viewer

Available Checkpoints

Base Models

Instruction-Tuned Models

Evaluation

Open-Ended LLM Benchmarks

Detoxification

Detoxification task performances are evaluated on the Monet-VD 1.4B model.

RealToxicityPrompts

ToxiGen

Examples

Text Generation

from transformers import pipeline

model_name = "MonetLLM/monet-vd-1.4B-100BT-hf"
pipe = pipeline(
    "text-generation",
    model_name,
    tokenizer=AutoTokenizer.from_pretrained(model_name),
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,
)
print(pipe("The key to life is", max_new_tokens=20, do_sample=True)[0]["generated_text"])

Code Generation

from transformers import pipeline

model_name = "MonetLLM/codemonet-vd-1.4B-100BT-hf"
pipe = pipeline(
    "text-generation",
    model_name,
    tokenizer=AutoTokenizer.from_pretrained(model_name),
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,
)

text = '''
def print_len(x: str):
    """For a given string x, print the length of x."""
'''
print(pipe(text, max_new_tokens=10)[0]["generated_text"].split("\n\n")[0])

Chat Completion

from transformers import pipeline

model_name = "MonetLLM/codemonet-vd-1.4B-100BT-chat-hf"
pipe = pipeline(
    "text-generation",
    model_name,
    tokenizer=AutoTokenizer.from_pretrained(model_name),
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,
)

text = tokenizer.apply_chat_template(
    [{"role": "user", "content": "Hi! How are you?"}],
    add_generation_prompt=True,
    tokenize=False,
)
print(pipe(text, max_new_tokens=30, do_sample=True)[0]["generated_text"])

Using vLLM

The custom implementation of vLLM is provided in the repository.

from vllm import LLM, ModelRegistry, SamplingParams
from modeling_monet_vllm import MonetForCausalLM

# Register Monet architecture with vLLM
ModelRegistry.register_model("MonetForCausalLM", MonetForCausalLM)

model = LLM(
    "MonetLLM/monet-vd-1.4B-100BT-hf",
    trust_remote_code=True,
    dtype="bfloat16",
    gpu_memory_utilization=0.8
)
sampling_params = SamplingParams(max_tokens=20, temperature=1.0)
print(model.generate("The key to life is", sampling_params)[0].outputs[0].text)

Training

Model

• Architecture: Monet
• Pretraining tokens: 100B
• Precision: bfloat16

Hardware

• TPUs: TPU-v4-64 Pod Slice (supported by TRC Program)

Software

• Training Framework: Jax, Flax

Intended Use

Primary Intended Uses

This model is designed to advance research on language models and serve as a foundational component for generative AI-driven functionalities. Its primary applications, mostly in English, include:

• Mechanistic interpretability research for language models
• Text generation with enhanced interpretability
• Code generation (CodeMonet variant)
• Chat completion (instruction-tuned variant)
• Vision-language tasks (VisionMonet variant)

Out-of-Scope Uses

This model has not been explicitly developed or tested for all potential downstream applications. Therefore:

1. Limitations & Mitigations: Developers should be mindful of common language model limitations, and thoroughly evaluate and mitigate risks regarding accuracy, safety, and fairness—especially in high-stakes or high-risk scenarios.
2. Legal & Regulatory Compliance: Developers must comply with any applicable laws and regulations (e.g., privacy, trade compliance), taking into account the model’s English-focused training (refer to FineWeb-Edu).
3. No License Modification: Nothing in this Model Card modifies or restricts the license under which this model is released.
4. Unsupported Programming Languages: Programming in languages not covered by StarCoderData(CodeMonet variant) is not within the model’s intended scope.

Model Architecture

Monet introduces a novel Mixture-of-Experts (MoE) architecture with several key innovations:

• Parameter-efficient expert decomposition: overall parameter count grows in proportion to the square root of the number of experts
• Fine-grained expert specialization: offers clear insight into model behavior
• Precise manipulation of knowledge: enables control over domain knowledge, programming language capabilities, and toxicity level.

Ethical Considerations

Transparency

• Designed specifically for enhanced interpretability
• Enables understanding of internal model behavior
• Allows tracking of knowledge attribution

Control

• Supports toxicity mitigation
• Enables domain-specific knowledge control
• Maintains performance while adjusting behavior

License and Usage

Monet is licensed under the Apache 2.0 license. The model is primarily intended for research and educational use. Important licensing notes:

• Instruction-tuned models have been fine-tuned using a dataset mix with outputs generated from third party models
• Research and educational use is encouraged
• Commercial use is subject to Apache 2.0 license terms

Citation

@article{park2024monet,
      title={{Monet: Mixture of Monosemantic Experts for Transformers}}, 
      author={Jungwoo Park and Young Jin Ahn and Kee-Eung Kim and Jaewoo Kang},
      journal={arXiv preprint arXiv:2404.05567},
      year={2024}
}