config.json

{
  "dataset": {
    "mask_probability": 0.15,
    "max_seq_len": 1024
  },
  "model": {
    "amygd_d_latent": 64,
    "amygd_temp": 0.1,
    "d_model": 1024,
    "dropout": 0.1,
    "ffn_factor": 4,
    "initializer_range": 0.02,
    "num_heads": 16,
    "num_layers": 3,
    "vocab_size": 30000
  },
  "training": {
    "batch_size": 4,
    "betas": [
      0.9,
      0.98
    ],
    "clip_grad_norm": 1.0,
    "eps": 1e-06,
    "gradient_accumulation_steps": 32,
    "learning_rate": 0.0001,
    "mixed_precision": "bf16",
    "warmup_steps_perc": 0.1,
    "weight_decay": 0.01
  }
}