mamba-790m-resoning

Mamba-790M Reasoning Faithfulness Model

Fine-tuned Mamba-790M model for evaluating reasoning trace faithfulness in language models.

Model Description

This is a LoRA-adapted version of Mamba-790M trained to generate responses that faithfully follow provided reasoning traces.

• Base Model: state-spaces/mamba-790m-hf (790M parameters)
• Adapter Type: LoRA (Low-Rank Adaptation)
• Training Data: Dolphin-DeepSeek Reasoning Dataset

Intended Use

This model is designed for research on:

• Reasoning faithfulness: Testing if model outputs align with stated reasoning
• AI interpretability: Understanding how models follow (or deviate from) reasoning traces
• Alignment research: Measuring consistency between reasoning and conclusions

Example Usage

from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import torch

# Load model
base_model = AutoModelForCausalLM.from_pretrained(
    "state-spaces/mamba-790m-hf",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

model = PeftModel.from_pretrained(base_model, "NakshJain/mamba-790m-resoning")
tokenizer = AutoTokenizer.from_pretrained("state-spaces/mamba-790m-hf", trust_remote_code=True)

# Format: <user>question</user><think>reasoning</think><answer>
prompt = "<user>What is 15 + 27?</user><think>Let me add: 15 + 27 = 42</think><answer>"

inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=100)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)

print(response)
# Output: 42

Training Details

Training Data

• Source: Dolphin-DeepSeek filtered ShareGPT conversations
• Training Set: 18500 reasoning examples
• Validation Set: 1500 examples
• Test Set: 500 examples
• Format: Structured as <user>query</user><think>reasoning</think><answer>response</answer>

Training Configuration

• LoRA Rank (r): 32
• LoRA Alpha: 32
• LoRA Dropout: 0.05
• Target Modules: in_proj, x_proj, dt_proj
• Learning Rate: 1.5e-5
• Batch Size: 4 (effective: 8 with gradient accumulation)
• Epochs: 1
• Optimizer: AdamW
• LR Schedule: Cosine with 3% warmup

Citation

If you use this model, please cite:

@misc{mamba-reasoning-faithfulness-2024,
  author = {Naksh Jain},
  title = {Mamba-790M Reasoning Faithfulness Model},
  year = {2024},
  publisher = {HuggingFace},
  url = {https://huggingface.co/NakshJain/mamba-790m-resoning}
}

Acknowledgments

• Base Model: Mamba by Tri Dao and Albert Gu
• Training Dataset: Dolphin-DeepSeek filtered by PJMixers-Dev
• Framework: HuggingFace Transformers, PEFT

License

Apache 2.0 (inherits from base Mamba model)