The diffusion model fails to converge during LoRA fine-tuning, and the loss keeps oscillating.

[krona12]

I am encountering a issue where the diffusion model's training loss fails to converge when fine-tuning with LoRA (system.use_lora=True). Instead of decreasing, the loss oscillates significantly throughout the training process.

To isolate the problem, I attempted to fine-tune the model using the minimal example dataset provided (the single object in ./data/shape_diffusion/objaverse). Using the configuration file detailed below, I ran the training for 200 epochs with the command: CUDA_VISIBLE_DEVICES=0 python train.py --config path/to/config.yaml --train --gpu 0 system.use_lora=True. The expected behavior of a steadily decreasing loss did not occur. Instead, the loss value remained unstable and oscillated for the entire duration, as shown in the attached loss curve.
I have confirmed this behavior occurs both when training with labels and without them; neither setup leads to convergence. Interestingly, despite the unstable training, the fine-tuned model checkpoint is able to produce seemingly normal inferences.

This behavior is not limited to the example data. My initial attempt was on a custom dataset built from parnet-mobility objects, which I processed using the data/watertight_and_sampling.py script. During that process, I first encountered a NaN loss, an issue similar to one reported by jseobyun. I resolved the NaN loss by filtering out the problematic sharp_surface data. However, even after fixing the data and retraining for about 100 epochs, the primary issue of the oscillating loss persisted.

Since the problem occurs on both a larger custom dataset and the minimal single-object example, I suspect that the issue may originate from the data, the training configuration, or perhaps even the code. I would appreciate any guidance or suggestions on how to resolve this convergence issue.

Loss curvs is as follows

Configuration (step1x-3d-geometry-label-1300m.yaml)

# [INFO] 
#   | Name             | Type                    | Params
# -------------------------------------------------------------
# 0 | shape_model      | MichelangeloAutoencoder | 191 M 
# 1 | visual_condition | Dinov2CLIPEncoder       | 731 M 
# 2 | label_condition  | LabelEncoder            | 8.2 K 
# 3 | denoiser_model   | FluxDenoiser            | 1.3 B
# -------------------------------------------------------------
# 1.3 B     Trainable params
# 923 M     Non-trainable params
# 2.2 B     Total params

exp_root_dir: "outputs_test"
name: "step1x-3d-geometry-label/dinov2reglarge518-clip-label-fluxflow-dit1300m"
#tag: "${rmspace:${system.shape_model_type}+n${data.n_samples}+${system.optimizer.name}lr${system.optimizer.args.lr},_}"
tag: "objaverse-200ep-wlabel-wcaption-lora-t7181228-testpret"
seed: 0

data_type: "Objaverse-datamodule"
data:
  root_dir: "./data/shape_diffusion/objaverse"
  load_geometry: True                # whether to load geometry
  geo_data_type: "sdf" 
  with_sharp_data: True
  n_samples: 32768
  noise_sigma: 0.
  random_flip: True # whether to randomly flip the input mesh
  random_color_jitter: True # whether to add random color jitter to the input images
  random_rotate: True # whether to add random rotation to the input images

  load_image: True                # whether to load images 
  image_type: "rgb"               # rgb, normal
  idx: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19]        # front view
  n_views: 1
  background_color: [255, 255, 255]

  load_caption: True             # whether to load captions #False
  load_label: True          # whether to load labels #True 

  batch_size: 8 # TODO: change this for your own dataset
  num_workers: 4 # change this for your own dataset

system_type: "rectified-flow-system"
system:
  val_samples_json: "val_data/images/val_samples_rgb_image.json"
  guidance_scale: 7.5
  num_inference_steps: 30
  eta: 0.0

  shape_model_type: michelangelo-autoencoder
  shape_model:
    pretrained_model_name_or_path: /home/yuantingyu/models
    subfolder: Step1X-3D-Geometry-Label-1300m
    n_samples: ${data.n_samples}
    with_sharp_data: ${data.with_sharp_data}
    use_downsample: true
    num_latents: 2048
    embed_dim: 64
    point_feats: 3
    out_dim: 1
    num_freqs: 8
    include_pi: false
    heads: 12
    width: 768
    num_encoder_layers: 8
    num_decoder_layers: 16
    use_ln_post: true
    init_scale: 0.25
    qkv_bias: false
    use_flash: true #true
    use_checkpoint: false

  visual_condition_type: "dinov2-clip-encoder"
  visual_condition:
    pretrained_dino_name_or_path: /home/yuantingyu/models/dinov2-large
    pretrained_clip_name_or_path: /home/yuantingyu/models/clip-vit-large-patch14
    encode_camera: false
    n_views: ${data.n_views}
    empty_embeds_ratio: 0.1
    normalize_embeds: false
    zero_uncond_embeds: true
    image_size: 224

  label_condition_type: "label-encoder"
  label_condition:
    hidden_size: 1024
    empty_embeds_ratio: 0.1
    normalize_embeds: false
    zero_uncond_embeds: true

  denoiser_model_type: "flux-denoiser"
  denoiser_model:
    pretrained_model_name_or_path: "/home/yuantingyu/models/modelscope/Step1X-3D-Geometry-Label-1300m/transformer/diffusion_pytorch_model.safetensors"
  
    input_channels: ${system.shape_model.embed_dim}
    width: 1536
    layers: 8
    num_single_layers: 16
    num_heads: 16
    use_visual_condition: True
    visual_condition_dim: 1024
    n_views: ${data.n_views}
    use_label_condition: True
    label_condition_dim: ${system.label_condition.hidden_size}

  noise_scheduler_type: "diffusers.schedulers.FlowMatchEulerDiscreteScheduler"
  noise_scheduler:
    num_train_timesteps: 1000
    
  denoise_scheduler_type: "diffusers.schedulers.FlowMatchEulerDiscreteScheduler"
  denoise_scheduler:
    num_train_timesteps: 1000

  loggers:
    wandb:
      enable: false
      project: "step1x-3d"
      name: image-to-shape-diffusion+${name}+${tag}

  loss:
    loss_type: "mse"
    lambda_diffusion: 1.

  optimizer:
    name: AdamW
    args:
      lr: 1.e-5                    
      betas: [0.9, 0.99]
      eps: 1.e-6

trainer:
  num_nodes: 1
  max_epochs: 200

  log_every_n_steps: 5 
  num_sanity_val_steps: 1
  val_check_interval: 1.0
  enable_progress_bar: true
  precision: 32                     
  strategy: 'deepspeed_stage_2'

checkpoint:
  save_last: true
  save_top_k: -1
  every_n_epochs: 10

[SerdarHelli]

I’m facing the same issue. Just to confirm — despite the unstable loss behavior, your fine-tuned model was still able to produce good results?

[krona12]

I’m facing the same issue. Just to confirm — despite the unstable loss behavior, your fine-tuned model was still able to produce good results?

---

You maybe should check your .yaml config and add an extra parameter pretrained_model_name_or_pathunder denoiser_model. The original config doesn’t include this by default, so even if you set system.use_lora=True, it actually doesn’t load the pretrained weights and instead trains from scratch. That’s why the loss behavior is the same with and without LoRA.

Even after adding this adjustment, in my case the loss still keeps oscillating — but the model can still generate normal outputs. I also compared the parameters before and after LoRA fine-tuning, and they did change (though only slightly), which confirms that LoRA did take effect. The model outputs are reasonable, but the loss doesn’t show a clear decreasing trend. For reference, I trained on the example dataset provided by the project.

Additionally, I’d like to share one more observation: I downloaded the same object (00a1a602456f4eb188b522d7ef19e81b) from Objaverse and processed it with the watertight_and_sampling.py script. However, the resulting .npz seems different from the one provided by the project. When I pass my .npz through the VAE encoder-decoder, the reconstruction quality is worse compared to using the .npz included in the project.

Below are some visual comparisons for reference:

• Raw 00a1a602456f4eb188b522d7ef19e81b.obj from Objaverse without texture:
  

• VAE reconstruction from .npz generated by running the sampling script:
  

• VAE reconstruction from the .npz provided in the project(./data/shape_diffusion/objaverse/surfaces/000-000/00a1a602456f4eb188b522d7ef19e81b.npz):
  

---

[SerdarHelli]

I successfully loaded the model and was able to train it. By the way, oscillation in the diffusion loss is expected, but in my case, it was excessive , it s not going down. Also, I wrote my own sampling script and was able to get good results from VAE. In my case, I’m working on generating a complete shape from a defected one, so I changed the conditioning input ,which comes from VAE, accordingly. However, when I applied this new condition, I couldn’t get meaningful loss behave either you , so I was stopping tranings , maybe I should go.

My recommendation for you , you may want to add gradient accumulation because your batch size is very small.

[mooncloud-max]

I successfully loaded the model and was able to train it. By the way, oscillation in the diffusion loss is expected, but in my case, it was excessive , it s not going down. Also, I wrote my own sampling script and was able to get good results from VAE. In my case, I’m working on generating a complete shape from a defected one, so I changed the conditioning input ,which comes from VAE, accordingly. However, when I applied this new condition, I couldn’t get meaningful loss behave either you , so I was stopping tranings , maybe I should go.

My recommendation for you , you may want to add gradient accumulation because your batch size is very small.

Hi, can you share the sampling script? My VAE reconstruction result is not good too.

[SerdarHelli]

I successfully loaded the model and was able to train it. By the way, oscillation in the diffusion loss is expected, but in my case, it was excessive , it s not going down. Also, I wrote my own sampling script and was able to get good results from VAE. In my case, I’m working on generating a complete shape from a defected one, so I changed the conditioning input ,which comes from VAE, accordingly. However, when I applied this new condition, I couldn’t get meaningful loss behave either you , so I was stopping tranings , maybe I should go.
My recommendation for you , you may want to add gradient accumulation because your batch size is very small.

Hi, can you share the sampling script? My VAE reconstruction result is not good too.

from pathlib import Path
from typing import Tuple, Union,Optional
import numpy as np
import open3d as o3d
import copy
from collections import defaultdict

import torch

def compute_dihedral_angle(mesh, edge_to_faces, edge):
    """Compute the dihedral angle between two triangles sharing an edge."""
    adjacent_faces = edge_to_faces[edge]
    if len(adjacent_faces) != 2:
        return 0  # Edge is not shared by exactly two faces

    # Compute normals of the two adjacent triangles
    normals = np.asarray(mesh.triangle_normals)
    normal1 = normals[adjacent_faces[0]]
    normal2 = normals[adjacent_faces[1]]

    # Calculate the dihedral angle
    cosine_angle = np.dot(normal1, normal2)
    angle = np.arccos(np.clip(cosine_angle, -1.0, 1.0))
    return angle

def extract_edges(mesh):
    """Extract all edges and their adjacent faces from a TriangleMesh."""
    triangles = np.asarray(mesh.triangles)
    edges_to_faces = defaultdict(list)

    for face_idx, triangle in enumerate(triangles):
        edges = [
            tuple(sorted((triangle[0], triangle[1]))),
            tuple(sorted((triangle[1], triangle[2]))),
            tuple(sorted((triangle[2], triangle[0]))),
        ]
        for edge in edges:
            edges_to_faces[edge].append(face_idx)

    return edges_to_faces

def sharp_edge_sampling(mesh, angle_threshold=np.pi / 18, sample_density=35):
    """
    Perform sharp edge sampling on a TriangleMesh.

    Args:
        mesh: Open3D TriangleMesh.
        angle_threshold: Dihedral angle threshold for detecting sharp edges.
        sample_density: Number of points to sample per unit length of edge.

    Returns:
        np.ndarray: Sampled points as a NumPy array.
    """
    edge_to_faces = extract_edges(mesh)

    # Identify sharp edges
    sharp_edges = []
    for edge in edge_to_faces:
        angle = compute_dihedral_angle(mesh, edge_to_faces, edge)
        if angle > angle_threshold:
            sharp_edges.append(edge)

    # Sample points along sharp edges
    vertices = np.asarray(mesh.vertices)
    sampled_points = []
    for edge in sharp_edges:
        point1 = vertices[edge[0]]
        point2 = vertices[edge[1]]
        edge_length = np.linalg.norm(point2 - point1)
        num_samples = max(2, int(edge_length * sample_density))
        for i in range(num_samples):
            t = i / (num_samples - 1)
            sampled_point = point1 * (1 - t) + point2 * t
            sampled_points.append(sampled_point)

    return np.array(sampled_points)

def sharp_edge_sampling_with_normals(mesh, angle_threshold=np.pi / 18, sample_density=35):
    """
    Perform sharp edge sampling on a TriangleMesh, including normals for sampled points.

    Args:
        mesh: Open3D TriangleMesh.
        angle_threshold: Dihedral angle threshold for detecting sharp edges.
        sample_density: Number of points to sample per unit length of edge.

    Returns:
        tuple: A tuple containing:
            - np.ndarray: Sampled points as a NumPy array.
            - np.ndarray: Normals of the sampled points as a NumPy array.
    """
    edge_to_faces = extract_edges(mesh)

    # Identify sharp edges
    sharp_edges = []
    for edge in edge_to_faces:
        angle = compute_dihedral_angle(mesh, edge_to_faces, edge)
        if angle > angle_threshold:
            sharp_edges.append(edge)

    # Sample points and compute normals along sharp edges
    vertices = np.asarray(mesh.vertices)
    vertex_normals = np.asarray(mesh.vertex_normals)
    sampled_points = []
    sampled_normals = []

    for edge in sharp_edges:
        point1 = vertices[edge[0]]
        point2 = vertices[edge[1]]
        normal1 = vertex_normals[edge[0]]
        normal2 = vertex_normals[edge[1]]

        edge_length = np.linalg.norm(point2 - point1)
        num_samples = max(2, int(edge_length * sample_density))

        for i in range(num_samples):
            t = i / (num_samples - 1)
            sampled_point = point1 * (1 - t) + point2 * t
            sampled_normal = normal1 * (1 - t) + normal2 * t
            normal_magnitude = np.linalg.norm(sampled_normal)

            if normal_magnitude > 1e-6:  # Avoid division by zero
                sampled_normal /= normal_magnitude
            else:
                sampled_normal = np.array([0.0, 0.0, 0.0])  # Default to zero vector if the magnitude is too small

            sampled_points.append(sampled_point)
            sampled_normals.append(sampled_normal)

    return np.array(sampled_points), np.array(sampled_normals)




def normalize_mesh_unitsphere(mesh: o3d.geometry.TriangleMesh) -> Tuple[o3d.geometry.TriangleMesh, np.ndarray, float]:
    """
    Normalize a 3D mesh: scale and center to fit within a unit sphere.

    Parameters:
        mesh: Open3D triangle mesh.

    Returns:
        normalized_mesh: Normalized Open3D triangle mesh.
        center: Center of the original mesh.
        scale: Scaling factor used for normalization.

    """
    mesh = copy.deepcopy(mesh)
    vertices = np.asarray(mesh.vertices)
    center = vertices.mean(axis=0)
    scale = np.linalg.norm(vertices - center, axis=1).max()  # Max distance
    scale = scale + (scale*0.01)  # add ofsset
    mesh.translate(-center)
    mesh.scale(1 / scale, center=(0, 0, 0))

    return mesh, center, scale



mesh = o3d.io.read_triangle_mesh("your?path")

mesh_o3d, center, scale = normalize_mesh_unitsphere(mesh)

mesh_o3d.compute_vertex_normals()

# Sample surface points uniform
pcd_o3d_uniform = mesh_o3d.sample_points_uniformly(number_of_points=100_000)
pcd_o3d_uniform.estimate_normals()
# RNG initialization
rng = np.random.default_rng()

# --- Uniform Sampling ---
pcd_uniform_np = np.asarray(pcd_o3d_uniform.points)
normals_uniform_np = np.asarray(pcd_o3d_uniform.normals)
idx_uniform = rng.choice(len(pcd_uniform_np), size=32768, replace=False)
pcd_uniform = pcd_uniform_np[idx_uniform]
normals_uniform = normals_uniform_np[idx_uniform]
uniform_data = np.concatenate([pcd_uniform, normals_uniform], axis=1)  # shape: [32768, 6]

# --- Sharp Edge Sampling ---
pcd_ses_np, normals_ses_np = sharp_edge_sampling_with_normals(
    mesh_o3d, angle_threshold=np.pi / 18, sample_density=35
)
pcd_ses_np = np.asarray(pcd_ses_np)
normals_ses_np = np.asarray(normals_ses_np)
replace_ses = len(pcd_ses_np) < 32768
idx_ses = rng.choice(len(pcd_ses_np), size=32768, replace=replace_ses)
pcd_ses = pcd_ses_np[idx_ses]
normals_ses = normals_ses_np[idx_ses]
sharp_data = np.concatenate([pcd_ses, normals_ses], axis=1)  # shape: [32768, 6]

# --- Convert to torch and send to GPU with batch dim ---
uniform_tensor = torch.tensor(uniform_data, dtype=torch.float32).unsqueeze(0).cuda()  # [1, 32768, 6]
sharp_tensor = torch.tensor(sharp_data, dtype=torch.float32).unsqueeze(0).cuda()

be ensure your model is watertight

[jseobyun]

I’ve been struggling with stable training as well, so I thought this might be helpful to share.

Aside from occasional NaN values in normals sampled from sharp edges, I also observed that some SDF values in volume_rand_points fall outside the expected [-1, 1] range. In my case, up to 1000–2000 such values were found in some samples. It might be worth checking the SDF value range in your training data.

For example, here’s a log from one of my samples:

[volume_rand_points] Invalid SDF values found:
  - NaN values: 0
  - Inf values: 0
  - Out of range (|sdf| > 1.0): 1781
  - Total invalid: 1781 / 100000
  - Valid SDF range: [-0.4334, 0.9997]
  - Valid SDF mean: 0.4271, std: 0.2808
Inspecting data: 309it [00:06, 47.75it/s]
Checking: 01_N3-90428.obj

[SerdarHelli]

I’ve been struggling with stable training as well, so I thought this might be helpful to share.

Aside from occasional NaN values in normals sampled from sharp edges, I also observed that some SDF values in volume_rand_points fall outside the expected [-1, 1] range. In my case, up to 1000–2000 such values were found in some samples. It might be worth checking the SDF value range in your training data.

For example, here’s a log from one of my samples:

[volume_rand_points] Invalid SDF values found:
  - NaN values: 0
  - Inf values: 0
  - Out of range (|sdf| > 1.0): 1781
  - Total invalid: 1781 / 100000
  - Valid SDF range: [-0.4334, 0.9997]
  - Valid SDF mean: 0.4271, std: 0.2808
Inspecting data: 309it [00:06, 47.75it/s]
Checking: 01_N3-90428.obj

Did this improve the training result ?

[jseobyun]

I’ve been struggling with stable training as well, so I thought this might be helpful to share.
Aside from occasional NaN values in normals sampled from sharp edges, I also observed that some SDF values in volume_rand_points fall outside the expected [-1, 1] range. In my case, up to 1000–2000 such values were found in some samples. It might be worth checking the SDF value range in your training data.
For example, here’s a log from one of my samples:

[volume_rand_points] Invalid SDF values found:
  - NaN values: 0
  - Inf values: 0
  - Out of range (|sdf| > 1.0): 1781
  - Total invalid: 1781 / 100000
  - Valid SDF range: [-0.4334, 0.9997]
  - Valid SDF mean: 0.4271, std: 0.2808
Inspecting data: 309it [00:06, 47.75it/s]
Checking: 01_N3-90428.obj

Did this improve the training result ?

The loss_logits begins to diverge after approximately 80–90 epochs of Shape VAE training in my case.

Since less than 1% of the 100,000 volume_rand_points exhibit this behavior, it does not significantly affect the average loss. However, when the sampling happened to be biased toward these noisy points, the loss spiked and failed to recover to a normal scale thereafter.

I think it will naturally cause performance degradation of diffusion parts too.