Would you please share some training details for stabilityai/stable-audio-open-small?

[yynil]

I'm trying to replicate the diffusion training process using benjamin-paine/freesound-laion-640k and benjamin-paine/free-music-archive-large.

Currently I use the same small network defined by stabilityai/stable-audio-open-small which is 16 layers with 1024 hidden size. I trained 3 epoches for now and continued training. But the demo reconstruction is not that good.
The diffusion mse loss is shake around 0.85.

The cfg-7 conditioned audio mel spectrum looks like:

Although it seemed learn to generate low frequency generation, the high frequency one is noisy.

Would you please share some details about the training data and the training time and loss trends? I'm not sure if I should terminate the training and add more data.

I just copied almost everything training details from stabilityai/stable-audio-open-small:

"diffusion": {
            "cross_attention_cond_ids": ["prompt", "seconds_total"],
            "global_cond_ids": ["seconds_total"],
            "diffusion_objective": "rectified_flow",
            "distribution_shift_options": {
                "min_length": 256,
                "max_length": 4096
            },
            "type": "dit",
            "config": {
                "io_channels": 64,
                "embed_dim": 1024,
                "depth": 16,
                "num_heads": 8,
                "cond_token_dim": 768,
                "global_cond_dim": 768,
                "transformer_type": "continuous_transformer",
                "attn_kwargs": {
                    "qk_norm": "ln"
                }
            }
        },
        "io_channels": 64
    },
    "training": {
        "use_ema": true,
        "log_loss_info": false,
        "pre_encoded": false,
        "timestep_sampler": "trunc_logit_normal",
        "optimizer_configs": {
            "diffusion": {
                "optimizer": {
                    "type": "AdamW",
                    "config": {
                        "lr": 2e-4,
                        "betas": [0.9, 0.95],
                        "eps": 1e-8,
                        "weight_decay": 0.01,
                        "foreach": true
                    }
                },
                "scheduler": {
                    "type": "InverseLR",
                    "config": {
                        "inv_gamma": 1000000,
                        "power": 0.5,
                        "warmup": 0.995
                    }
                }
            }
        },
        "demo": {
            "demo_every": 2000,
            "demo_steps": 50,
            "num_demos": 8,
            "demo_cond": [
                {"prompt": "Amen break 174 BPM", "seconds_total": 6},
                {"prompt": "People talking in a crowded cafe", "seconds_total": 10},
                {"prompt": "A short, beautiful piano riff in C minor", "seconds_total": 6},
                {"prompt": "Tight Snare Drum", "seconds_total": 1},
                {"prompt": "A dog barking next to a waterfall", "seconds_total": 6},
                {"prompt": "Glitchy bass design, I used Serum for this", "seconds_total": 4},
                {"prompt": "Synth pluck arp with reverb and delay, 128 BPM", "seconds_total": 6},
                {"prompt": "Birds singing in the forest", "seconds_total": 10}
            ],
            "demo_cfg_scales": [1, 4, 7]
        }
    }

[zqevans]

The loss doesn't seem unreasonable, are you starting from the pre-trained SAO Small base model in the repo? And with the pre-trained VAE? And what batch size?

And have you modified the demo prompts to be more reasonable for the prompt structure and seconds_total for your training dataset?

[yynil]

I load the weights from stabilityai/stable-audio-open-small, and then freeze vae part, replace the diffusion part's architect, then only update the diffusion model since I wanna try other architect other than transformer.
The Laion-640K's median length should be around 30 seconds and the free-music's median length should be longer than 30 seconds. I'm now analyzing the dataset, I guess the dataset should be selected with the training targets( around 11 seconds). I'll try to select the proper audios and report the training results

[zqevans]

If you're using your own architectures, I can't really speak to the training speed or quality. Training a diffusion model from scratch can take a long time depending on your dataset size, batch size, and compute budget.

If you're starting from our base RF model, you'll want to use base_model.ckpt checkpointin the SAO Small Hugging Face repo. That's the base rectified flow model, as opposed to the ARC post-trained model, which is model.ckpt in that repo.

[yynil]

I found the essential mistakes I made that I load the arc model. I'll correct it with the base_model.ckpt and run the training process again with the filtered data.

Now the training batch size is 16 with 4x4090, the dataset is around 100K valid audios with bfloat16 mixing precision. I also add fp8 training supports with deepspeed and torchao which enables the batch size to 32. After validating the configurations, I will move the codes to a fp8 enabled server to run a larger training.

Thanks a lot, I'll report the results later.

[sskalnik]

Any luck?

Looks like you have been working on training the ARC model as well?

I see that you are using the following in your fork:
precision = "bf16-mixed"
You may get better results if you use the default of 16-mixed.

I recommend avoiding gradient accumulation unless you must use it for a specific reason.