Update README.md

README.md

@@ -1,60 +1,126 @@
 ---
 license: apache-2.0
+language:
+  - en
 tags:
   - text-generation
   - causal-lm
-  - cosmicfish
-  - hrm
   - adaptive-reasoning
+  - hierarchical-reasoning
+  - hrm
   - custom-architecture
-language:
-  - en
+  - compact-model
+datasets:
+  - HuggingFaceFW/fineweb
+  - wikipedia
+  - openwebtext
+  - c4
+arxiv: 2605.28919
 ---
 
 # CosmicFish-HRM
 
-CosmicFish-HRM is a compact 82.77M parameter causal language model built around a Hierarchical Reasoning Module (HRM) that dynamically allocates reasoning compute during inference. Developed at Mistyoz AI.
+**Paper:** [CosmicFish-HRM: Adaptive Reasoning via Hierarchical Recurrent Mechanisms in Compact Language Models](https://arxiv.org/abs/2605.28919)
 
-Rather than applying a fixed number of transformer layers to every input, CosmicFish-HRM iterates through high-level and low-level reasoning cycles and uses a learned halting head to decide when to stop. Harder inputs trigger deeper reasoning trajectories while simpler ones halt early.
+**GitHub:** [MistyozAI/CosmicFish-HRM](https://github.com/MistyozAI/CosmicFish-HRM)
+
+CosmicFish-HRM is a compact 82.77M parameter causal language model built around a Hierarchical Reasoning Module (HRM) that dynamically allocates reasoning compute during inference. Rather than applying a fixed number of forward-pass layers to every input, the model iterates through high-level and low-level reasoning cycles and uses a learned halting head to decide when to stop. Harder inputs trigger deeper reasoning trajectories while simpler ones halt early.
+
+Built at Mistyoz AI, Hyderabad.
+
+---
 
 ## Architecture
 
+![Architecture](architecture.png)
+
 ```
 Input Blocks (Transformer) -> HRM Core (H + L levels, variable steps) -> Output Blocks (Transformer) -> LM Head
 ```
 
-| Hyperparameter | Value |
+The HRM core maintains two interacting recurrent states operating at different abstraction levels. The high-level module captures slower, more abstract reasoning while the low-level module handles finer-grained local computation. After each reasoning step a lightweight halting head decides whether to continue or stop, conditioned on the mean-pooled high-level state.
+
+**Key components:**
+
+- Grouped-Query Attention (GQA) with 8 query heads and 4 KV heads
+- Rotary Positional Embeddings (RoPE)
+- SwiGLU feedforward layers
+- RMSNorm (pre-norm for I/O blocks, post-norm inside HRM)
+- Learned halt/continue Q-head controlling per-input reasoning depth
+- Step penalty in the training loss encouraging efficient halting
+
+## Model Specs
+
+| Parameter | Value |
 |---|---|
-| Parameters | 82.77M |
+| Total parameters | 82.77M |
 | Embedding dimension | 448 |
 | Vocabulary size | 50,304 |
 | Context length | 512 |
-| Input layers | 6 |
-| Output layers | 6 |
-| Attention heads | 8 (4 KV, GQA) |
+| Input transformer layers | 6 |
+| Output transformer layers | 6 |
 | HRM H-layers | 4 |
 | HRM L-layers | 4 |
 | Max HRM steps | 16 |
+| Attention heads | 8 (4 KV, GQA) |
 
-**Key components:**
-- Grouped-Query Attention (GQA) with RoPE
-- SwiGLU feedforward layers
-- RMSNorm (pre-norm for I/O blocks, post-norm inside HRM)
-- Learned halt/continue Q-head controlling per-input reasoning depth
-- Step penalty in training loss encouraging efficient halting
+## Evaluation
+
+Zero-shot benchmark results:
+
+| Model | HellaSwag | PIQA | WinoGrande |
+|---|---|---|---|
+| CosmicFish-HRM (82M) | 26.2 | 58.1 | 50.7 |
+| GPT-2 Small (117M) | 29.7 | 62.5 | 50.7 |
+| OPT-125M | 30.6 | 62.6 | 52.9 |
+| Pythia-160M | 29.4 | 62.1 | 52.8 |
+
+At compact scale a portion of the parameter budget is allocated to the HRM reasoning infrastructure rather than raw language modeling capacity, which accounts for the gap versus fixed-depth baselines of similar size. The paper argues this tradeoff becomes more favorable as model scale increases.
+
+## Adaptive Reasoning Behavior
+
+The primary contribution of CosmicFish-HRM is not benchmark accuracy but adaptive compute allocation. The model uses different numbers of reasoning steps depending on input complexity:
+
+| Prompt | Mean HRM Steps |
+|---|---|
+| "The capital of France is" | 2.78 |
+| "Photosynthesis is the process by which plants" | 4.77 |
+| "If all roses are flowers and some flowers fade quickly..." | 7.03 |
+| "A bat and a ball cost $1.10 in total..." | 8.40 |
+
+Average steps across benchmarks stay well below the 16-step maximum, with high variance across samples, confirming the halting mechanism is input-sensitive rather than collapsing to a fixed depth.
+
+| Benchmark | Mean Steps | Std Dev |
+|---|---|---|
+| HellaSwag | 3.03 | 6.26 |
+| PIQA | 1.87 | 5.13 |
+| WinoGrande | 0.95 | 3.78 |
+| Overall | 2.68 | 5.95 |
 
 ## Usage
 
-This model uses a custom architecture and requires `trust_remote_code=True`.
+This model uses a custom architecture. The model code is included in this repo as `modeling_hrm_cosmicfish.py`.
+
+**Standalone chat script (downloads automatically):**
+
+```bash
+pip install torch safetensors huggingface-hub transformers termcolor
+python chat_hf.py
+```
+
+**Load manually:**
 
 ```python
 import torch
 import json
 import tiktoken
 from safetensors.torch import load_file
+from huggingface_hub import snapshot_download
 from modeling_hrm_cosmicfish import HRMCosmicFish, HRMCosmicFishConfig
 
-with open("config.json") as f:
+cache_dir = snapshot_download("MistyozAI/CosmicFish-HRM")
+
+with open(f"{cache_dir}/config.json") as f:
     cfg = json.load(f)
 
 config = HRMCosmicFishConfig(
@@ -73,38 +139,22 @@ config = HRMCosmicFishConfig(
     dropout=0.0,
 )
 
-state_dict = load_file("model.safetensors")
+state_dict = load_file(f"{cache_dir}/model.safetensors")
 model = HRMCosmicFish(config)
 model.load_state_dict(state_dict)
 model.eval()
 
 tokenizer = tiktoken.get_encoding("gpt2")
-
 prompt = "Artificial intelligence is"
 tokens = tokenizer.encode(prompt)
 idx = torch.tensor(tokens, dtype=torch.long).unsqueeze(0)
 
 with torch.no_grad():
-    output = model.generate(idx, max_new_tokens=50, temperature=0.7, top_k=40)
+    output = model.generate(idx, max_new_tokens=100, temperature=0.7, top_k=40)
 
 print(tokenizer.decode(output[0].tolist()))
 ```
 
-## Training
-
-CosmicFish-HRM was trained on the 10B-token CosmicSet dataset spanning web text, Wikipedia, code, mathematics, and research papers. Training used cosine LR decay with linear warmup, bfloat16 mixed precision, and gradient clipping.
-
-## Citation
-
-```bibtex
-@misc{cosmicfish-hrm,
-  title={CosmicFish-HRM: Adaptive Reasoning via Hierarchical Recurrent Mechanisms in Compact Language Models},
-  author={Venkat Akhil Lakkapragada},
-  year={2026},
-  howpublished={\url{https://huggingface.co/MistyozAI/CosmicFish-HRM}}
-}
-```
-
 ---
 
-Mistyoz AI, Hyderabad
+Mistyoz AI, Hyderabad