vae/chained_flow_vae_config.json

{
  "model_args": {
    "vae_type": "transformer_hidden",
    "hidden_size": 4096,
    "latent_size": 1024,
    "intermediate_size": 3072,
    "num_layers": 2,
    "num_heads": 8,
    "max_sequence_length": 16,
    "dropout": 0.0,
    "device": "cuda"
  },
  "data_args": {
    "dataset_path": "data/flow_cache/vae9b_3k_fp16",
    "dataset_split": "train",
    "tokens_per_epoch": 65536,
    "token_seed": 0,
    "response_only": true,
    "validation_fraction": 0.1,
    "split_seed": 0,
    "sequence_length": 8
  },
  "loss_args": {
    "lambda_mse": 1.0,
    "lambda_cos": 0.2,
    "lambda_norm": 0.0001,
    "beta": 3e-05,
    "free_bits": 0.0
  }
}