{
  "datasets": [
    {
      "name": "chatalpaca_predictive_state",
      "repo_id": "BRlkl/chatalpaca-multiturn-predictive-state3",
      "source_split": "train",
      "format": "predictive_state_messages",
      "validation_ratio": 0.02,
      "split_seed": 17,
      "min_turns": 2,
      "max_turns": 6,
      "max_message_chars": 6000,
      "use_base_chat_template": true
    }
  ],
  "model": {
    "base_model_name": "google/t5gemma-l-l-prefixlm-it",
    "initial_model_path": "BRlkl/test",
    "dtype": "bfloat16",
    "attn_implementation": "sdpa",
    "magicnorm_eps": 1e-06,
    "z_slots": 256,
    "num_time_tokens": 0,
    "use_explicit_time_features": false,
    "gate_attention_heads": 4,
    "max_observation_tokens": 1024,
    "max_decoder_tokens": 1024,
    "thought_loop_proposal_mode": "observation_hidden_compression",
    "preserve_observation_encoder_manifold": true,
    "observation_encoder_use_state_context": true,
    "latent_attention_mask_mode": "full",
    "initial_update_gate_bias": 20.0
  },
  "training": {
    "seed": 17,
    "num_workers": 2,
    "gradient_checkpointing": true,
    "mixed_precision": "bf16",
    "max_grad_norm": 1.0,
    "weight_decay": 0.01,
    "backbone_learning_rate": 5e-06,
    "new_module_learning_rate": 0.0001,
    "adam_beta1": 0.9,
    "adam_beta2": 0.95,
    "adam_epsilon": 1e-08,
    "fused_adamw": true,
    "freeze_gate_head": true,
    "log_every_steps": 1,
    "eval_every_steps": 50,
    "checkpoint_every_steps": 5000,
    "eval_max_batches": 16,
    "max_train_examples": 10,
    "max_validation_examples": 10,
    "freeze_decoder_first_fraction": 0.0,
    "response_loss_weight": 0.45,
    "probe_loss_weight": 0.35,
    "probe_hard_loss_weight": 0.15,
    "noop_probe_loss_weight": 0.05,
    "probe_hard_top_fraction": 0.2,
    "max_probe_branches_per_step": 8,
    "max_probe_noop_ticks": 4,
    "require_probe_steps": 1,
    "drop_llm_not_stated_probes": true,
    "drop_llm_positive_answer_not_in_prefix": true,
    "drop_llm_future_answer_leaks": true,
    "validation_behavior_enabled": true,
    "validation_behavior_exact_only": true,
    "validation_behavior_max_batches": 4,
    "validation_behavior_max_probe_generations": 192,
    "validation_probe_max_new_tokens": 96,
    "validation_probe_extra_new_tokens": 8,
    "validation_noop_survival_threshold_ticks": 64,
    "validation_self_fed_behavior_enabled": true,
    "validation_behavior_self_fed_max_batches": 2,
    "validation_behavior_self_fed_max_probe_generations": 96,
    "wandb_train_metric_keys": [
      "train/loss_total",
      "train/loss_probe",
      "train/probe_first_token_exact_match"
    ],
    "wandb_validation_metric_keys": [
      "loss_total",
      "validation_goal_loss",
      "behavior/probe_exact_match",
      "behavior/self_fed/probe_exact_match",
      "behavior/noop_survival_exact",
      "behavior/secret_word_exact",
      "behavior/entity_binding_exact",
      "behavior/correction_latest_state_accuracy",
      "behavior/not_stated_accuracy"
    ],
    "checkpoint_selection_metric": "validation_goal_loss",
    "checkpoint_selection_mode": "min",
    "validation_probe_miss_penalty": 2.0,
    "validation_noop_miss_penalty": 1.0,
    "validation_self_fed_miss_penalty": 1.0,
    "validation_not_stated_miss_penalty": 0.5,
    "validation_entity_miss_penalty": 0.75,
    "validation_correction_miss_penalty": 0.75,
    "checkpoint_min_probe_exact_match": 0.5,
    "checkpoint_min_secret_word_exact": 0.5,
    "checkpoint_min_entity_binding_exact": 0.5,
    "checkpoint_min_not_stated_accuracy": 0.5,
    "checkpoint_min_correction_latest_state_accuracy": 0.5,
    "checkpoint_min_noop_survival_exact": 0.5,
    "checkpoint_min_self_fed_probe_accuracy": 0.4,
    "self_feedback_start_ratio": 0.0,
    "self_feedback_end_ratio": 0.15,
    "self_feedback_start_fraction": 0.65,
    "self_feedback_end_fraction": 0.95,
    "self_feedback_seed": 17,
    "self_feedback_max_new_tokens": 160,
    "self_feedback_extra_new_tokens": 8,
    "student_observation_seed": 17,
    "student_prev_assistant_dropout_prob": 0.0,
    "student_prev_assistant_keep_last_chars": null
  },
  "phase": {
    "micro_batch_size": 16,
    "eval_batch_size": 16,
    "gradient_accumulation_steps": 1,
    "num_train_epochs": 100,
    "warmup_ratio": 0.03,
    "shuffle_train": true
  },
  "cache": {
    "preprocessed_root": "cache/preprocessed_pre_sft_multiturn_predictive"
  },
  "paths": {
    "run_root": "runs_pre_sft_multiturn_predictive",
    "export_root": "exports_multiturn_predictive"
  },
  "inference": {
    "format": "predictive_state_multiturn",
    "use_base_chat_template": true
  },
  "wandb": {
    "enabled": true,
    "project": "samantha-pre-sft",
    "run_name": "t5gemma2-thoughtloop-pre-sft-predictive-state"
  },
  "hub": {
    "model_repo_id": "BRlkl/test_multiturn_predictive_state_overfit",
    "private": false
  }
}