Add TRM-Nano README with architecture details and usage"

trm_solver/trm_nano/README.md

@@ -0,0 +1,89 @@
+# TRM-Nano: Tiny Recursive Model for NeuroGolf 2026
+
+Based on [Samsung TRM](https://arxiv.org/abs/2510.04871) (ARC Prize 2025 1st Place Paper Award).
+
+## Architecture
+
+| Parameter | Value | Notes |
+|-----------|-------|-------|
+| hidden_size | 96 | 24 per head |
+| num_heads | 4 | non-causal attention |
+| num_layers | 2 | single block, weight-shared |
+| expansion | 4 | SwiGLU FFN |
+| H_cycles (T) | 3 | outer recursions |
+| L_cycles (n) | 4 | inner recursions per T |
+| Total depth | 15 passes | T×(n+1) through same 2-layer block |
+| Params | ~200K | ~0.8 MB FP32 ONNX |
+| puzzle_emb | 8 tokens | learned per-task |
+
+## Key Design Decisions
+
+1. **Weight sharing**: Single 2-layer block recursed 15 times. This is the core TRM insight — recursive depth > width.
+2. **Unrolled recursion**: No `Loop`/`Scan` ops in ONNX. All 15 passes hardcoded sequentially.
+3. **Post-norm residual**: `x = rms_norm(x + attn(x))` then `x = rms_norm(x + mlp(x))`.
+4. **Dual latent**: `z_H` (answer) and `z_L` (reasoning). L-cycles update reasoning with input context, then one H-update refines the answer.
+5. **Input injection**: `z_L` receives `z_H + x_embed` as injection. `z_H` receives `z_L` only (no input). This forces `z_H` to represent a complete answer.
+
+## NeuroGolf I/O
+
+- **Input**: `[1, 10, 30, 30]` float32 one-hot → argmax → flat 900 tokens (+2 shift)
+- **Output**: logits → softmax over color channels → `[1, 10, 30, 30]`
+- **ONNX size**: ~800 KB (well within 1.44 MB limit)
+
+## Usage
+
+### Single task training
+```bash
+cd trm_solver/trm_nano
+python train_task.py \
+    --task_file /path/to/task.json \
+    --output_dir onnx_out/ \
+    --epochs 500 \
+    --hidden_size 96 \
+    --device cuda
+```
+
+### From Kaggle notebook
+```python
+import sys
+sys.path.insert(0, '/kaggle/input/arc-agi-code/trm_solver/trm_nano')
+
+from model import TRMNano, TRMNanoConfig
+from train_task import train_on_task, evaluate_task, export_onnx
+
+config = TRMNanoConfig(hidden_size=96, H_cycles=3, L_cycles=4)
+model = train_on_task(task_data, config, epochs=500, device="cuda")
+export_onnx(model, "task_id.onnx")
+```
+
+## Training Details
+
+- **Augmentation**: 8 dihedral transforms × random color permutation (keep 0 fixed) × random translation = ~1000 augmented views per task
+- **Optimizer**: AdamW, lr=1e-3, weight_decay=0.1, betas=(0.9, 0.95)
+- **Schedule**: Cosine annealing
+- **EMA**: decay=0.999 (final weights use EMA)
+- **Loss**: Cross-entropy on output tokens, ignore PAD positions
+- **Evaluation**: Majority voting over 64 augmented forward passes with inverse transform
+
+## Comparison to Full TRM
+
+| | Full TRM | TRM-Nano |
+|---|---|---|
+| hidden_size | 512 | 96 |
+| params | 7M | ~200K |
+| T×n | 3×6=18+3=21 | 3×4=12+3=15 |
+| training | 3 days 4×H100 | ~5 min/task single GPU |
+| ARC-1 score | 44.6% (all tasks) | TBD (per-task) |
+| ONNX size | 26 MB | ~0.8 MB |
+
+## Files
+
+| File | Purpose |
+|------|---------|
+| `model.py` | TRM-Nano architecture (PyTorch) |
+| `train_task.py` | Per-task training + eval + ONNX export |
+
+## References
+
+- [Less is More: Recursive Reasoning with Tiny Networks](https://arxiv.org/abs/2510.04871)
+- [Samsung-TRM code](https://huggingface.co/wtfmahe/Samsung-TRM)