如何确保神经微分方程训练的收敛性？Julia语言SciML教程练习6.3疑难求解

Hey there! Let's work through your problem with replacing the Lotka-Volterra dy/dt equation with a neural network. The issues you're seeing—initial loss explosions and stuck convergence—are super common in regression tasks with neural networks, especially when dealing with physical systems where input/output scales can vary. Here are actionable steps to fix both:

1. Normalize Your Data (Critical First Step) #

Neural networks are extremely sensitive to input and target scales. If your u (x, y) values or du[2] targets have large ranges, random initial weights can produce outputs that are way off, leading to those 10^20 loss values.

How to implement: #

First, compute the mean and standard deviation (or min/max) of your training data, then normalize both inputs and targets to a standard range (like mean 0, std 1). After training, you can reverse this normalization to get predictions in the original scale.

using Statistics

# Assume u_data is a matrix where each column is a (x,y) sample
# du2_target is a vector of true dy/dt values
u_mean = mean(u_data, dims=2)
u_std = std(u_data, dims=2)
du2_mean = mean(du2_target)
du2_std = std(du2_target)

# Normalize data
u_norm = (u_data .- u_mean) ./ u_std
du2_norm = (du2_target .- du2_mean) ./ du2_std

# Reverse normalization for predictions
function denorm_du2(pred_norm)
    return pred_norm .* du2_std .+ du2_mean
end

You can also integrate normalization directly into your network with Flux's Normalize layer for cleaner code:

NN = Chain(
    Normalize(2; u_mean, u_std),  # Normalize input (x,y)
    Dense(2, 30, relu),
    Dense(30, 1),
    x -> x .* du2_std .+ du2_mean  # Denormalize output back to original scale
)

2. Use Robust Neural Network Initialization #

The default weight initialization for Flux's Dense layers might not be ideal for your regression task. Using schemes tailored to your activation function can prevent extreme initial outputs.

Recommendations: #

• For relu activations, use He initialization (Flux.kaiming_uniform) which accounts for relu's tendency to zero out half the neurons.
• For smoother activations like swish, Xavier initialization (Flux.glorot_uniform) works well.

# Example with He initialization for relu layers
NN = Chain(
    Normalize(2; u_mean, u_std),
    Dense(2, 30, relu; init=Flux.kaiming_uniform),
    Dense(30, 1; init=Flux.kaiming_uniform),
    x -> x .* du2_std .+ du2_mean
)

You can also initialize biases to small values (like 0.1 instead of 0) to avoid starting with all outputs near zero, which can slow early training.

3. Switch to a More Robust Loss Function #

Mean Squared Error (MSE) penalizes large errors heavily, which can cause loss explosions if initial outputs are way off. Huber Loss is a great alternative—it acts like MSE for small errors but like MAE (Mean Absolute Error) for large errors, making it more robust to outliers and bad initializations.

using Flux: huber_loss

function loss_fn(NN, u, target)
    pred = NN(u)
    return huber_loss(pred, target)
end

4. Tune Your Optimizer & Add Learning Rate Scheduling #

ADAM is a great optimizer, but the default learning rate (0.001) might be too high for your task, leading to noisy updates that prevent late-stage convergence. Adding learning rate scheduling helps the optimizer fine-tune parameters as training progresses.

Try these adjustments: #

• Lower the initial learning rate to 0.0005 or even 0.0001.
• Use a cosine annealing scheduler to gradually reduce the learning rate over iterations, which helps the optimizer settle into a good minimum.

using Flux: Optimiser, ADAM, CosineAnnealing

# Set up optimizer with learning rate scheduling
base_opt = ADAM(0.0005)
scheduler = CosineAnnealing(1000)  # Reduce LR over 1000 iterations
opt = Optimiser(scheduler, base_opt)

5. Add Gradient Clipping to Prevent Explosions #

If you still see gradient explosions (even with normalization), gradient clipping limits the maximum norm of the gradients, ensuring parameter updates don't get too large.

# Add gradient clipping to your optimizer chain
opt = Optimiser(ClipNorm(1.0), scheduler, base_opt)

6. Adjust Network Structure for Better Convergence #

Your current 2->30->1 network is simple, but small tweaks can help:

• Add a BatchNorm layer after the first dense layer to stabilize training by normalizing activations between layers:

  NN = Chain(
      Normalize(2; u_mean, u_std),
      Dense(2, 30),
      BatchNorm(30, relu),  # BatchNorm + activation
      Dense(30, 1),
      x -> x .* du2_std .+ du2_mean
  )
  

• Try a smoother activation function like swish instead of relu—it has no hard zero point, which can lead to more stable gradients:

  NN = Chain(
      Normalize(2; u_mean, u_std),
      Dense(2, 30, swish; init=Flux.glorot_uniform),
      Dense(30, 1; init=Flux.glorot_uniform),
      x -> x .* du2_std .+ du2_mean
  )
  

7. Use Mini-Batch Training Instead of Full Batch #

If you're training on the entire dataset at once, switching to mini-batches can introduce beneficial noise that helps the optimizer escape local minima, leading to better convergence.

using Flux: DataLoader

# Create a DataLoader with batch size 32 (adjust based on your dataset size)
loader = DataLoader((u_norm, du2_norm), batchsize=32, shuffle=true)

# Train using mini-batches
for epoch in 1:2000
    Flux.train!(loss_fn, Flux.params(NN), loader, opt)
    # Print loss every 100 epochs to monitor progress
    if epoch % 100 == 0
        current_loss = loss_fn(NN, u_norm, du2_norm)
        println("Epoch $epoch, Loss: $current_loss")
    end
end

Final Tips #

• Always split your data into training and validation sets—this lets you check if your model is overfitting (if training loss keeps dropping but validation loss plateaus). Use early stopping to halt training when validation loss stops improving.
• Start with a smaller network (e.g., 2->16->1) before scaling up—smaller networks are easier to train and less prone to overfitting.

内容的提问来源于stack exchange，提问作者Bruce Kuwahara