train with continuous node feature

[yusuke0324]

Hi @xinyangATK,

I'm currently training a model on a custom dataset where node features consist of both continuous values and multi-categorical attributes (represented as multi-hot vectors).

From my understanding, when using continuous features, the get_logit_beta_stats_con function should be used instead of get_logit_beta_stats in the precondition module.

However, I noticed that the implementation of the expected value of logit(z_t) in the code seems to differ from the equation described in the paper (Appendix B, Eq. (29)). Specifically, the current code uses:

E1 = 1.0 / (eta * alpha_t * xmin) * (
    (eta * alpha_t * xmax).lgamma() - (eta * alpha_t * xmin).lgamma())
Shouldn't the denominator be eta * alpha_t * (xmax - xmin) instead of eta * alpha_t * xmin?

It seems that E2, V1, and V2 may have the same issue as well.

After applying this correction, I observed that the values in X within sample_batch start to diverge with each iteration, and the model eventually returns NaNs.
I suspect this may be related to the behavior of get_logit_beta_stats_con, or possibly how scale_shift, sigmoid_start, sigmoid_end, and sigmoid_power are defined.

Could you kindly advise on this?

[anastasiiakim]

Also interested in this. @yusuke0324 Did you manage to solve it?

[yusuke0324]

I ended up not using precondition and it still successfully worked.
Currently working on trying to clamp min max after precondition to prevent the issue.

[xinyangATK]

Hi @yusuke0324 ,

Thank you very much for your attention. You're absolutely right, the denominator should be eta * alpha_t * (xmax - xmin). Since I didn't use the continuous precondition function in original version, I didn't have too many checks. Really thanks for telling us this issue, I will update the right function.

However, the updated continuous precondition function works well for my project. I also suspect that some hpyerparameters cause the NaNs in your case, could you mind share the scale / shape of your custom data and your own hpyerparameters? I'll try to have a look.

btw, please looking forward our new project about beta diffusion, the precondition function maybe drop in the new model which will be easy to apply beta diffusion.

[yusuke0324]

Thank you for your testing. The following setting is used for my train. It seems like the mean is going to be negative value and std to under 1 as t increases, then input value becomes larger than 64 bits.

dataset:
    name: 'x'
    remove_h: null
    datadir: 'data/rplan/'
    rand_perm: False
    feat:
        type:
            - deg
            - eig1
            - eig2
        scale: 10
        norm: True
    use_small_dataset: False

general:
    name: 'x-graph-tf-model'
    forward_visualization: False
    sample_visualization: False
    vis_number: 10
    wandb: 'online'
    wandb_project: 'graphgen'
    gpus: [0, ]
    resume: null
    test_only:
    check_val_every_n_epochs: 50
    sample_every_val: 1
    val_check_interval: null
    samples_to_generate: 200
    samples_to_save: 20
    chains_to_save: 1
    log_every_steps: 50
    number_chain_steps: 50
    final_model_samples_to_generate: 300
    final_model_samples_to_save: 20
    final_model_chains_to_save: 20
    evaluate_all_checkpoints: False

train:
    n_epochs: 2000
    batch_size: 128
    lr: 0.001
    clip_grad: 1.0
    save_model: True
    num_workers: 12
    ema_decay: 0.999
    progress_bar: false
    weight_decay: 1e-12
    optimizer: adamw
    seed: 0

model:
    type: 'discrete'
    transition: 'marginal'
    model: 'graph_tf'
    diffusion_steps: 1000
    n_layers: 6
    eta_from: 'train'
    extra_features:
    high_order: False
    hidden_mlp_dims: {'X': 256, 'E': 128, 'y': 128}
    hidden_dims : {'dx': 256, 'de': 64, 'dy': 64, 'n_head': 8, 'dim_ffX': 256, 'dim_ffE': 128, 'dim_ffy': 128}
    lambda_train: [5, 0]
    sigmoid_start: 10
    sigmoid_end: -13
    sigmoid_power: 1.
    scale_shift:
        node: [0.9, 0.09]
        edge: [0.9, 0.09]
    eta:
        node: [1000, 100, 100, 1]
        edge: [1000, 100, 100, 1]
    input_space: 'logit'
    pre_condition: True
    noise_feat_type: 'x'
    default_min_max: [-1., 1.]
    max_feat_num: [52, 16, 1]
    cibcentration_m: False
    concentration_strategy: 'deg'
    use_alignment_loss: true
    alignment_sigma: 0.01
    alignment_threshold: 0.05
    alignment_gamma: 3
    lambda_max_manhattan: 5.0
    lambda_max_colinear: 0.5
    lambda_max_orthogonal: 0.5

hydra:
  job:
    chdir: True
  run:
    dir: ./outputs/${dataset.name}/${now:%Y-%m-%d}/${now:%H-%M-%S}-${general.name}

[xinyangATK]

I think the eta should be at least larger than 1 due to some properties in beta distribution. For the smallest eta, I suggest you set it as 10 and you can try it later.

[yusuke0324]

Thank you for the suggestion. I set [1000, 100, 100, 10] for eta but still got NaN. During the reverse process, X_t becomes larger and ended up overflow. In fact, how much the precondition impacts?

I also tried clamp after precondition to make sure the preconditioned value is within the original range. However, the KLUB didn't decrease with clamp during training.

[yusuke0324]

This is a toy reverse process code that loops 1000 steps and apply precondition over t.
I realized model.forward and sampling (t->t-1) doesn't change X value mutch (at least in the beginning train epoch ~50 epochs).
The X_t with eta=1000 becomes inf around t=500.

import torch
import matplotlib.pyplot as plt
from src.gbd_utils.precondition import PreConditionMoudle

# ---- setting ----
device = torch.device("cpu")
dataset_name = "rplan"
Shift = {"node": 0.09, "edge": 0.09}
Scale = {"node": 0.9, "edge": 0.9}
sigmoid_start = 10
sigmoid_end = -13
sigmoid_power = 1.0
eta_list = [1000, 100, 10]
input_space = 'logit'
type_name = 'node'

# ---- precond module setup ----
precon = PreConditionMoudle(
    dataset_name=dataset_name,
    dataset_info=None,
    Shift=Shift,
    Scale=Scale,
)
precon.device = device

# ---- alpha_t schedule ----
T = 1000
sample_steps = 1 - torch.arange(T) / (T - 1) * (1. - 1e-5)
sample_steps = torch.cat([sample_steps, torch.zeros(1)])
logit_alpha = sigmoid_start + (sigmoid_end - sigmoid_start) * (sample_steps ** sigmoid_power)
alpha_t_vals = torch.sigmoid(logit_alpha)

# ---- init X ----
X0 = 0.01 * torch.ones(1, 1, 1, device=device)
X0 = precon.scale_shift(X0, type='node')

# ---- log ----
traces = {}

# ---- toy reverse process ----
with torch.no_grad():
    for eta_val in eta_list:
        eta_tensor = torch.full((1, 1, 1), eta_val, device=device)
        X_t = X0.clone()
        mean_trace, std_trace, X_trace, alpha_trace = [], [], [], []

        for alpha_val in alpha_t_vals:
            alpha_tensor = torch.full((1, 1, 1), alpha_val.item(), device=device)
            mean_t, std_t = precon.pre_condition_fn(
                alpha_t=alpha_tensor,
                eta=eta_tensor,
                type=type_name,
                prior_prob=None,
                input_space=input_space
            )
            X_t = (X_t - mean_t) / std_t
            mean_trace.append(mean_t.squeeze().item())
            std_trace.append(std_t.squeeze().item())
            X_trace.append(abs(X_t.item()) + 1e-10)  # for log display
            alpha_trace.append(alpha_val.item())

        traces[eta_val] = {
            "mean": mean_trace,
            "std": std_trace,
            "X": X_trace,
            "alpha": alpha_trace
        }

# ---- plot ----
plt.figure(figsize=(12, 8))
for eta_val in eta_list:
    plt.plot(traces[eta_val]["mean"], label=f"eta={eta_val}")
plt.title("E[logit(z_t)] over diffusion steps")
plt.xlabel("t")
plt.ylabel("E[logit(z_t)]")
plt.ylim(-10, 10)
plt.grid(True)
plt.legend()

plt.figure(figsize=(12, 8))
for eta_val in eta_list:
    plt.plot(traces[eta_val]["std"], label=f"eta={eta_val}")
plt.title("Std[logit(z_t)] over diffusion steps")
plt.xlabel("t")
plt.ylabel("Std[logit(z_t)]")
plt.ylim(-10, 10)
plt.grid(True)
plt.legend()

plt.figure(figsize=(12, 8))
for eta_val in eta_list:
    plt.plot(traces[eta_val]["X"], label=f"eta={eta_val}")
plt.yscale("log")
plt.title("X_t over diffusion steps (log scale)")
plt.xlabel("t")
plt.ylabel("X_t (log scale)")
plt.grid(True)
plt.legend()

plt.figure(figsize=(12, 8))
plt.plot(alpha_t_vals.tolist(), label="alpha_t")
plt.title("alpha_t schedule")
plt.xlabel("t")
plt.ylabel("alpha_t")
plt.grid(True)
plt.legend()

plt.tight_layout()
plt.show()
 

In training steps, model's inputs are always within specific range so model must not learn such a large scale input?

[xinyangATK]

I think the x_t before precondition should be in logit space, but your example shows the x_t is in original space. For x_t in logit space and original space, their precondition functions are different and our precondition function provided in this repo is in logit space, so you may try to forward x_t in logit space via apply_noise function.

btw, your reverse process seems not right because you can not get next x_t through precondition on x_t+1. You should use the network to predict current state. I think you can visualize the forward process without network. Our precondition is applied on noise data like x_t, but x_t can not obtained by simply normalizing (applying precondition on) x_t-1, there should be a beta diffusion schedule for this forward process.

[yusuke0324]

Right, I forgot to change x_t to logit space before the loop but even logit x_t goes to inf after applying precondition many times.
Please note that the code is a toy code that imitates sample_batch to see x_t changes over t after applying precondition.
I omitted network prediction because the prediction can't output a large enough value to prevent inf after converting to logit. I realized log_delta's range in sample_new_logit_x is [-90, 0] because of clamp within torch.log such as torch.log(torch.tensor(MIN)) is approximately -87. In this case, output X_s by X_s = self.sample_new_logit_x(X_s, X, alpha_t, alpha_s, eta_x) is very similar to arg X_s.

I think it occurs due to while log_delta in sample_new_logit_x is clamped around -90, the precondition can increase infinity.