By JamesSand
Here is the code for training the second order model.
Please cite our related papers which using second order method.
High-Order Matching for One-Step Shortcut Diffusion Models
@article{chen2025high,
title={High-Order Matching for One-Step Shortcut Diffusion Models},
author={Chen, Bo and Gong, Chengyue and Li, Xiaoyu and Liang, Yingyu and Sha, Zhizhou and Shi, Zhenmei and Song, Zhao and Wan, Mingda},
journal={arXiv preprint arXiv:2502.00688},
year={2025}
}Hofar: High-order augmentation of flow autoregressive transformers
@article{liang2025hofar,
title={Hofar: High-order augmentation of flow autoregressive transformers},
author={Liang, Yingyu and Sha, Zhizhou and Shi, Zhenmei and Song, Zhao and Wan, Mingda},
journal={arXiv preprint arXiv:2503.08032},
year={2025}
}@article{cao2025force,
title={Force Matching with Relativistic Constraints: A Physics-Inspired Approach to Stable and Efficient Generative Modeling},
author={Cao, Yang and Chen, Bo and Li, Xiaoyu and Liang, Yingyu and Sha, Zhizhou and Shi, Zhenmei and Song, Zhao and Wan, Mingda},
journal={arXiv preprint arXiv:2502.08150},
year={2025}
}
The First Order model is defined at here
class MLP(nn.Module):
def __init__(self, input_dim=2, hidden_num=100):
super().__init__()
self.fc1 = nn.Linear(input_dim+1, hidden_num, bias=True)
self.fc2 = nn.Linear(hidden_num, hidden_num, bias=True)
self.fc3 = nn.Linear(hidden_num, input_dim, bias=True)
self.act = lambda x: torch.tanh(x)
def forward(self, x_input, t):
# We only input zero order here, so we have a dimension d input
inputs = torch.cat([x_input, t], dim=1)
x = self.fc1(inputs)
x = self.act(x)
x = self.fc2(x)
x = self.act(x)
x = self.fc3(x)
return xThe Second Order model is defined at here
class MLP_2nd_order(nn.Module):
def __init__(self, input_dim=2, hidden_num=100):
super().__init__()
self.fc1 = nn.Linear(input_dim + input_dim + 1, hidden_num, bias=True)
self.fc2 = nn.Linear(hidden_num, hidden_num, bias=True)
self.fc3 = nn.Linear(hidden_num, input_dim, bias=True)
self.act = lambda x: torch.tanh(x)
def forward(self, first_order_input, x_input, t):
# We concate zero order and first order here, to get a dimension 2d input
inputs = torch.cat([first_order_input, x_input, t], dim=1)
x = self.fc1(inputs)
x = self.act(x)
x = self.fc2(x)
x = self.act(x)
x = self.fc3(x)
return xfirst order weight: 1; second order weight: 1e-11
first order weight: 1; second order weight: 1e-8
first order weight: 1e6; second order weight: 1
pip install -r requirements.txtpython second_order_code.pyPlease refer to model_eval.ipynb
