Upload TRMForCausalLM

config.json

@@ -10,32 +10,30 @@
     "H_layers": 0,
     "batch_size": 1,
     "causal": true,
-    "dropout": 0.1,
+    "dropout": 0.2,
     "expansion": 4,
     "forward_dtype": "bfloat16",
     "freeze_embeddings": false,
-    "full_grad_recursion": true,
     "halt_exploration_prob": 0.1,
     "halt_max_steps": 8,
     "hidden_size": 768,
     "loss": {
-      "loss_type": "softmax_cross_entropy",
+      "loss_type": "stablemax_cross_entropy",
       "name": "losses@ACTLossHead"
     },
-    "mlp_t": false,
     "name": "recursive_reasoning.trm@TinyRecursiveModel",
     "no_ACT_continue": true,
-    "num_heads": 8,
-    "num_layers": 4,
+    "num_heads": 12,
+    "num_layers": 8,
     "num_puzzle_identifiers": 1,
     "pos_encodings": "rope",
     "pretrained_embeddings_model": "Alibaba-NLP/gte-modernbert-base",
     "puzzle_emb_len": 0,
     "puzzle_emb_ndim": 0,
-    "residual_recursion": true,
-    "seq_len": 256,
+    "residual_recursion": false,
+    "seq_len": 512,
     "vocab_size": 50368,
     "y_cycles": 2,
-    "z_cycles": 2
+    "z_cycles": 3
   }
 }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:cf1b29f3e306475a0f20b9ea682611acabb21f017ce4e20fc4e348162d9ed862
-size 422719008
+oid sha256:a90414327a4218c51b6ad370d10fdb3719d49fb9c7ddc7750b42f119ec5cfadf
+size 535967176