Upload deepseek-tiny-mla-o-v0.1 model weights and documentation

README.md

@@ -0,0 +1,104 @@
+---
+license: apache-2.0
+language:
+- en
+datasets:
+- wikitext
+- glue
+pipeline_tag: text-generation
+tags:
+- transformer
+- attention
+- mla
+- research
+- output-subspace
+---
+
+# Deepseek Tiny Mla O V0.1
+
+6-layer DeepSeek-V3 with MLA + shared output latent space trained for research on shared subspaces in Transformer attention mechanisms.
+
+## Model Description
+
+- **Model Type**: Transformer Decoder (DeepSeek-V3 based)
+- **Architecture**: 6-layer decoder with Mixture of Experts
+- **Parameters**: 16.17M
+- **Hidden Size**: 256
+- **Attention Heads**: 8
+- **Head Dimension**: 32
+- **Sequence Length**: 1,024 tokens
+- **Query Latent Dimension**: 96
+- **Key-Value Latent Dimension**: 64
+- **Output Latent Dimension**: 96
+
+## Performance
+
+- **SST-2 Accuracy**: 86.24%
+- **WikiText-103 Perplexity**: 29.33
+
+## Research Context
+
+This model is part of the [shared-subspaces](https://github.com/chrisjmccormick/shared-subspaces) research project investigating the impact of shared output latent spaces in Transformer attention mechanisms.
+
+### Output Subspace Decomposition
+This model implements a shared output latent space where the attention output projection W^O is decomposed into:
+```
+W^O = W^OB · W^OA
+```
+Where W^OA are per-head projections to the latent space and W^OB is a shared projection back to the model dimension.
+
+## Usage
+
+```python
+import torch
+from transformers import DeepseekV3ForCausalLM, AutoTokenizer
+
+# Load model and tokenizer
+model = DeepseekV3ForCausalLM.from_pretrained("ChrisMcCormick/deepseek-tiny-mla-o-v0.1")
+tokenizer = AutoTokenizer.from_pretrained("ChrisMcCormick/deepseek-tiny-mla-o-v0.1")
+
+# For MLA-o model, apply the output subspace patch
+from patch_o_proj import patch_o_proj_implementation
+patch_o_proj_implementation(
+    model=model,
+    o_latent_dim=96,
+    variant="sequential_norm"
+)
+
+# Generate text
+inputs = tokenizer("The future of AI is", return_tensors="pt")
+outputs = model.generate(**inputs, max_length=50, temperature=0.7)
+print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+```
+
+## Training Details
+
+- **Pre-training Dataset**: WikiText-103
+- **Fine-tuning Dataset**: SST-2 (GLUE)
+- **Optimizer**: AdamW
+- **Learning Rate**: 5e-4 (pre-training), 5e-5 (fine-tuning)
+- **Weight Decay**: 0.01 (pre-training), 0.05 (fine-tuning)
+- **Precision**: bfloat16
+- **Compilation**: torch.compile with inductor backend
+- **Training Steps**: 12,500 (pre-training), 1,500 (fine-tuning)
+
+## Limitations
+
+- Small scale model (16M parameters) intended for research purposes
+- Trained on limited data compared to production models
+- May require custom loading code for output subspace variants
+
+## Citation
+
+```bibtex
+@misc{mccormick2025sharedsubspaces,
+  title={Shared Subspaces in Transformer Attention: Investigating Output Latent Spaces},
+  author={McCormick, Chris},
+  year={2025},
+  howpublished={\url{https://github.com/chrisjmccormick/shared-subspaces}}
+}
+```
+
+## License
+
+Apache 2.0

config.json

@@ -0,0 +1,53 @@
+{
+  "architectures": [
+    "DeepseekV3ForCausalLM"
+  ],
+  "attention_backend": "flash_attention_2",
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 50256,
+  "classifier_dropout": null,
+  "dtype": "float32",
+  "eos_token_id": 50256,
+  "ep_size": 1,
+  "first_k_dense_replace": 1,
+  "head_dim": 32,
+  "hidden_act": "silu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 256,
+  "initializer_range": 0.02,
+  "intermediate_size": 1024,
+  "kv_lora_rank": 64,
+  "max_position_embeddings": 1024,
+  "model_type": "deepseek_v3",
+  "moe_intermediate_size": 128,
+  "moe_layer_freq": 2,
+  "n_group": 1,
+  "n_routed_experts": 4,
+  "n_shared_experts": 1,
+  "norm_topk_prob": true,
+  "num_attention_heads": 8,
+  "num_experts_per_tok": 2,
+  "num_hidden_layers": 6,
+  "num_key_value_heads": 8,
+  "num_nextn_predict_layers": 1,
+  "pad_token_id": 50256,
+  "pretraining_tp": 1,
+  "q_lora_rank": 96,
+  "qk_head_dim": 32,
+  "qk_nope_head_dim": 0,
+  "qk_rope_head_dim": 32,
+  "rms_norm_eps": 1e-06,
+  "rope_interleave": true,
+  "rope_scaling": null,
+  "rope_theta": 10000.0,
+  "routed_scaling_factor": 1,
+  "scoring_func": "softmax",
+  "tie_word_embeddings": true,
+  "topk_group": 1,
+  "topk_method": "softmax_aux",
+  "transformers_version": "4.56.0",
+  "use_cache": false,
+  "v_head_dim": 32,
+  "vocab_size": 50257
+}

example_usage.py

@@ -0,0 +1,53 @@
+#!/usr/bin/env python3
+"""
+Example usage for deepseek-tiny-mla-o-v0.1
+"""
+
+import torch
+from transformers import DeepseekV3ForCausalLM, AutoTokenizer
+
+def main():
+    # Load model and tokenizer
+    print("Loading model...")
+    model = DeepseekV3ForCausalLM.from_pretrained("ChrisMcCormick/deepseek-tiny-mla-o-v0.1")
+    tokenizer = AutoTokenizer.from_pretrained("ChrisMcCormick/deepseek-tiny-mla-o-v0.1")
+    
+    # Apply output subspace patch for MLA-o model
+    print("Applying output subspace patch...")
+    from patch_o_proj import patch_o_proj_implementation
+    patch_o_proj_implementation(
+        model=model,
+        o_latent_dim=96,
+        variant="sequential_norm"
+    )
+    
+    # Set to evaluation mode
+    model.eval()
+    
+    # Example prompts
+    prompts = [
+        "The future of artificial intelligence is",
+        "In a world where technology advances rapidly,",
+        "The most important discovery in science was",
+    ]
+    
+    print("\nGenerating text...")
+    for prompt in prompts:
+        inputs = tokenizer(prompt, return_tensors="pt")
+        
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_length=50,
+                temperature=0.7,
+                do_sample=True,
+                pad_token_id=tokenizer.eos_token_id
+            )
+        
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        print(f"Prompt: {prompt}")
+        print(f"Generated: {generated_text}")
+        print("-" * 50)
+
+if __name__ == "__main__":
+    main()

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 50256,
+  "eos_token_id": [
+    50256
+  ],
+  "pad_token_id": 50256,
+  "transformers_version": "4.56.0",
+  "use_cache": false
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8e5c85a72d7fb506cf28f5fb2d5f2a9539f60b9d49ffcc46c65d3d1803b45823
+size 67839424

special_tokens_map.json

@@ -0,0 +1,6 @@
+{
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "pad_token": "<|endoftext|>",
+  "unk_token": "<|endoftext|>"
+}

tokenizer_config.json

@@ -0,0 +1,21 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "50256": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "model_max_length": 1024,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}