Optim

VMCOptimizer

vmc/optim/optimizer/VMCOptimizer

 1class VMCOptimizer(BaseVMCOptimizer):
 2
 3    def __init__(
 4        self,
 5        nqs: DDP,
 6        sampler_param: dict,
 7        electron_info: ElectronInfo,
 8        opt: Optimizer,
 9        lr_scheduler: Union[List[LRScheduler], LRScheduler] = None,
10        max_iter: int = 2000,
11        dtype: Dtype = None,
12        external_model: any = None,
13        checkpoint: str = None,
14        read_model_only: bool = False,
15        only_sample: bool = False,
16        pre_CI: CIWavefunction = None,
17        pre_train_info: dict = None,
18        clean_opt_state: bool = False,
19        noise_lambda: float = 0.05,
20        sr_config: SRConfig | None = None,
21        lm_config: LMConfig | None = None,
22        rgn_config: RGNConfig | None = None,
23        interval: int = 100,
24        prefix: str = "VMC",
25        MAX_AD_DIM: int = -1,
26        kfac: KFACPreconditioner | None = None,
27        use_clip_grad: bool = False,
28        max_grad_norm: float = 1.0,
29        max_grad_value: float = 1.0,
30        start_clip_grad: int = None,
31        clip_grad_method: str = "l2",
32        clip_grad_scheduler: Optional[Callable[[int], float]] = None,
33        use_3sigma: bool = False,
34        k_step_clip: int = 100,
35        use_spin_raising: bool = False,
36        spin_raising_coeff: float = 1.0,
37        only_output_spin_raising: bool = False,
38        spin_raising_scheduler: Optional[Callable[[int], float]] = None,
39    )

opt-params

 1from utils import ElectronInfo, Dtype
 2
 3opt_type = optim.AdamW
 4opt_params = {"lr": 0.001, "betas": (0.9, 0.999)}
 5opt = opt_type(model.parameters(), **opt_params)
 6
 7prefix = "vmc"
 8def clip_grad_scheduler(step):
 9   if step <= 4000:
10      max_grad = 1.0
11   elif step <= 8000:
12      max_grad = 0.1
13   else:
14      max_grad = 0.01
15   return max_grad
16
17vmc_opt_params = {
18    "nqs": model,
19    "opt": opt,
20    # "lr_scheduler": lr_scheduler,
21    # "read_model_only": True,
22    "dtype": dtype,
23    "sampler_param": sampler_param,
24    # "only_sample": True,
25    "electron_info": electron_info,
26    # "use_spin_raising": True,
27    # "spin_raising_coeff": 1.0,
28    # "only_output_spin_raising": True,
29    "max_iter": 5000,
30    "interval": 100,
31    "MAX_AD_DIM": 80000,
32    # "checkpoint": f"./h50/focus-init/checkpoint/H50-2.00-oao-mps-rnn-dcut-30-222-focus-20w-checkpoint.pth",
33    "prefix": prefix,
34    "use_clip_grad": True,
35    "max_grad_norm": 1,
36    "start_clip_grad": -1,
37    "clip_grad_scheduler": clip_grad_scheduler,
38}

  • nqs: Ansatz(e.g. Transformer, MPS-RNN, Graph-MPS-RNN).

  • opt: Optimizer(e.g., Adam, Adamw, SGD).

  • lr_scheduler: LRScheduler, Default: None.

  • read_model_only: Read model from the checkpoint file.

  • dtype: data-dtype: (e.g., Dtype(dtype=torch.complex128, device="cuda"))

  • sampler_param: see sample-param

  • only_sample: No calculating gradient. This is used to calculate energy.

  • max_iter: the number of the iteration.

  • interval: the time of the saving the checkpoint file.

  • MAX_AD_DIM: the nbatch of the backward.

  • checkpoint: Read model/optimizer/lr_scheduler from the checkpoint file, Default: None. The legacy keyword check_point is still accepted with a deprecation warning.

  • prefix: the prefix of the checkpoint file, e.g., vmc-checkpoint.pth.

  • use_clip_grad: clip gradient, Default: False.

  • max_grad_norm: the max of the l2-norm when clipping gradient.

  • start_clip_grad: clip gradient from the k-th iteration.

  • clip_grad_scheduler: the scheduler of clipping gradient, this is Callable[[int], float].

  • sr_config: configure SR/minSR. damping_lambda can be either a positive constant or a callable Callable[[int], float] receiving the optimization step. The default is the constant 1.0e-4.

 1from pynqs.optim import SRConfig
 2
 3# constant damping
 4sr_config = SRConfig(sr_method="minsr", damping_lambda=1.0e-4)
 5
 6# scheduled damping
 7sr_config = SRConfig(
 8    sr_method="minsr",
 9    damping_lambda=lambda step: max(1.0e-4 * 0.95**step, 1.0e-6),
10)

  • lm_config: configure the Linear method through LMConfig. The delta field can be either a non-negative constant or a callable Callable[[int], float] receiving the optimization step. With a constant value, PyNQS keeps the default schedule max(delta * 0.9**step, 1e-6).

1from pynqs.optim import LMConfig
2
3lm_config = LMConfig(delta=0.1)
4lm_config = LMConfig(delta=lambda step: max(0.1 * 0.9**step, 1.0e-6))

  • rgn_config: configure RGN through RGNConfig. The epsilon, delta, and damping_lambda fields can be constants or callables Callable[[int], float] receiving the optimization step.

1from pynqs.optim import RGNConfig
2
3rgn_config = RGNConfig(
4    epsilon=1.0,
5    delta=0.0,
6    damping_lambda=1.0e-3,
7)

Penalty coefficient schedules

Penaltywavefunctions accepts fixed coefficients or callables in alpha and beta. alpha(epoch) returns a symmetric (K, K) matrix with zero diagonal; beta(epoch) returns a scalar. Each schedule is evaluated at epoch zero during construction, then before each penalty VMC step using the same optimization-step index as SR damping. The resulting coefficients are shared by the loss, ordinary gradient, SR and minSR for that step.

alpha0 = torch.tensor([[0.0, 0.1], [0.1, 0.0]], dtype=torch.float64)

def alpha_schedule(epoch):
    return alpha0 * (1.0 + min(epoch / 1000, 1.0))

def beta_schedule(epoch):
    return 0.1 * min(epoch / 1000, 1.0)

ansatz = Penaltywavefunctions(
    single_ansatz=single_ansatz,
    alpha=alpha_schedule,
    beta=beta_schedule,
    nqubits=sorb,
    nele=nele,
    device=device,
)

Either coefficient may remain fixed. Schedules must be deterministic functions of the step, identical on every rank. Outputs are detached buffers, not trainable parameters. Checkpoints save the current coefficient tensors, not the functions; supply the schedules again when reconstructing the model. For a nonzero spin penalty, enable use_spin_raising=True in the optimizer so that the spin local quantities are evaluated, even if beta(0) is zero.

Rayleigh-Gauss-Newton

pynqs.optim.grad.rgn.RGN_grad implements the regularized Rayleigh-Gauss-Newton (RGN) update described by Peng and Chan, Phys. Rev. Research 7, 043351 (2025). It follows the same stochastic objects used by pynqs.optim.grad.lm.LM_grad:

\[O_i(n) = \frac{1}{\Psi(n)}\frac{\partial \Psi(n)}{\partial \theta_i}, \qquad h_i(n) = \partial_i E_{\rm loc}(n) + O_i(n)E_{\rm loc}(n).\]

The sampled gradient, overlap, and approximate Hessian are

\[\begin{split}G_i &= \operatorname{Cov}(E_{\rm loc}, O_i),\\ S_{ij} &= \operatorname{Cov}(O_i, O_j),\\ H^{\rm eff}_{ij} &= \operatorname{Cov}(O_i, h_j) - G_i\langle O_j\rangle - \langle E_{\rm loc}\rangle S_{ij}.\end{split}\]

RGN minimizes the second-order expansion with the SR overlap penalty, which gives the linear equation

\[\left(H^{\rm eff} + \delta I + \frac{S+\delta' I}{\epsilon}\right)c = -G.\]

In PyNQS the optimizer stores the preconditioned gradient

\[d\theta = \left(H^{\rm eff} + \delta I + \frac{S+\delta' I}{\epsilon}\right)^{-1}G,\]

and the optimizer step applies theta <- theta - lr*dtheta. For the paper’s convention use lr=1.

Usage in VMCOptimizer:

 1from pynqs.optim import RGNConfig
 2
 3vmc_opt = VMCOptimizer(
 4    nqs=model,
 5    opt=opt,
 6    sampler_param=sampler_param,
 7    electron_info=electron_info,
 8    rgn_config=RGNConfig(
 9        epsilon=1.0,
10        delta=0.0,
11        damping_lambda=1.0e-3,
12    ),
13)

RGNConfig.epsilon is the overlap-penalty scale \(\epsilon\). Small values approach SR, while math.inf gives the approximate Newton limit based on \(H^{\rm eff}\). RGNConfig.delta is the H-side shift, analogous to the \(\delta I\) added to L in LM. RGNConfig.damping_lambda is the S-side shift, analogous to the S-side shift added to R in LM and to damping_lambda in SR/minSR. For finite \(\epsilon\), damping_lambda enters the final matrix as \(\delta_s/\epsilon\). Each field can also be a Callable[[int], float] scheduler receiving the optimization step.

Optimizer

Linear method

Linear method ref.

    1. Chem. Phys. 152, 024111 (2020); doi: 10.1063/1.5125803

  • PHYSICAL REVIEW RESEARCH 7, 043351 (2025)

Linear method的梯度计算在 pynqs/optim/grad/lm.py 的函数 LM_grad 中， 欲使用之，只需向 VMCOptimizer 传入 lm_config 即可（例如 lm_config=LMConfig(...)）。 如果不给该字段传值则不会启用 LM。 除此之外，还有超参 \(\delta\) 需要调整，对应 LMConfig.delta，默认是 \(0.1\)。如果 delta 是常数，则按照 delta = max(delta * 0.9**(epoch), 1e-6) 进行衰减。 如果 delta 是 Callable[[int], float]，则每一步直接使用 delta(epoch) 的返回值作为 \(\delta\)。 理论上在优化最后，这项应该衰减至 \(0\).

该段代码包含计算梯度和更新两部分，计算梯度按照 J. Chem. Phys. 152, 024111 (2020) 中的方式实现，具体公式推导见文档。 简而言之，最后是构造一下广义本征值问题(GEVP)并求解：

\[\begin{split}Lc = \tilde{E} Rc,\quad L = \begin{bmatrix} E & G^\top_{\rm r} \\ G_{\rm c} & H \end{bmatrix} + \delta I , R = \begin{bmatrix} 1&\\ &S \end{bmatrix}+\delta'I\end{split}\]

with

\[\begin{split}[G_{\rm r}]_i &= \langle{h_i(n)}\rangle - E\langle{g_i(n)}\rangle\\ [G_{\rm c}]_i &= \sum_n p(n) \bar{\epsilon}(n) O_i(n)\\ S_{ij} &= \sum_n p(n) \bar{o}_i(n)O_j(n)\\ H_{ij} &= \sum_n p(n)\bar{o}_i(n) h_j(n) - [G_{\rm c}]_i\langle{O_j(n)}\rangle\end{split}\]

where

\[\begin{split}\bar{\epsilon}(n) &= E_{\rm loc}(n) - \langle{E_{\rm loc}(n)}\rangle\\ \bar{o}_i(n) &= O_i(n) - \langle{O_i(n)}\rangle\end{split}\]

and

\[h_i(n) = \partial_i E_{\rm loc}(n) + O_i(n)E_{\rm loc}(n),\quad \partial_i E_{\rm loc}(n) = \frac{\partial}{\partial\theta_i}\sum_{m\in SD}H_{nm}\frac{\varPsi(m)}{\varPsi(n)}\]
\[\begin{split}g_i(n) &= \frac{\langle n|\varPsi_i\rangle}{\langle n|\varPsi\rangle} = \frac{1}{\varPsi(n)} \frac{\partial\varPsi(n)}{\partial\theta_i} = O_i(n) \label{eq:gi}\\ h_i(n) &= \frac{\langle n|\hat{H}|\varPsi_i\rangle}{\langle n|\varPsi\rangle} = \partial_i E_{\rm loc}(n) + O_i(n)E_{\rm loc}(n)\end{split}\]

这里 \(|\varPsi_i\rangle = \partial_{\theta_i}|\varPsi\rangle\), 在更新的时候，实现了以上文章中类似线搜索的方式，见 try_step_update 函数。