Optim

VMCOptimizer

vmc/optim/optimizer/VMCOptimizer

class VMCOptimizer(BaseVMCOptimizer):

    def __init__(
        self,
        nqs: DDP,
        sampler_param: dict,
        electron_info: ElectronInfo,
        opt: Optimizer,
        lr_scheduler: Union[List[LRScheduler], LRScheduler] = None,
        max_iter: int = 2000,
        dtype: Dtype = None,
        external_model: any = None,
        checkpoint: str = None,
        read_model_only: bool = False,
        only_sample: bool = False,
        pre_CI: CIWavefunction = None,
        pre_train_info: dict = None,
        clean_opt_state: bool = False,
        noise_lambda: float = 0.05,
        sr_config: SRConfig | None = None,
        lm_config: LMConfig | None = None,
        rgn_config: RGNConfig | None = None,
        interval: int = 100,
        prefix: str = "VMC",
        MAX_AD_DIM: int = -1,
        kfac: KFACPreconditioner | None = None,
        use_clip_grad: bool = False,
        max_grad_norm: float = 1.0,
        max_grad_value: float = 1.0,
        start_clip_grad: int = None,
        clip_grad_method: str = "l2",
        clip_grad_scheduler: Optional[Callable[[int], float]] = None,
        use_3sigma: bool = False,
        k_step_clip: int = 100,
        use_spin_raising: bool = False,
        spin_raising_coeff: float = 1.0,
        only_output_spin_raising: bool = False,
        spin_raising_scheduler: Optional[Callable[[int], float]] = None,
    )

opt-params

from utils import ElectronInfo, Dtype

opt_type = optim.AdamW
opt_params = {"lr": 0.001, "betas": (0.9, 0.999)}
opt = opt_type(model.parameters(), **opt_params)

prefix = "vmc"
def clip_grad_scheduler(step):
   if step <= 4000:
      max_grad = 1.0
   elif step <= 8000:
      max_grad = 0.1
   else:
      max_grad = 0.01
   return max_grad

vmc_opt_params = {
    "nqs": model,
    "opt": opt,
    # "lr_scheduler": lr_scheduler,
    # "read_model_only": True,
    "dtype": dtype,
    "sampler_param": sampler_param,
    # "only_sample": True,
    "electron_info": electron_info,
    # "use_spin_raising": True,
    # "spin_raising_coeff": 1.0,
    # "only_output_spin_raising": True,
    "max_iter": 5000,
    "interval": 100,
    "MAX_AD_DIM": 80000,
    # "checkpoint": f"./h50/focus-init/checkpoint/H50-2.00-oao-mps-rnn-dcut-30-222-focus-20w-checkpoint.pth",
    "prefix": prefix,
    "use_clip_grad": True,
    "max_grad_norm": 1,
    "start_clip_grad": -1,
    "clip_grad_scheduler": clip_grad_scheduler,
}

• 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 can be either a positive constant or a callable Callable[[int], float] receiving the optimization step. The default is the constant 1.0e-4.

from pynqs.optim import SRConfig

# constant damping
sr_config = SRConfig(space="sample", damping=1.0e-4)

# scheduled damping
sr_config = SRConfig(
    space="sample",
    damping=lambda step: max(1.0e-4 * 0.95**step, 1.0e-6),
)

• 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).

from pynqs.optim import LMConfig

lm_config = LMConfig(delta=0.1)
lm_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.

from pynqs.optim import RGNConfig

rgn_config = RGNConfig(
    epsilon=1.0,
    delta=0.0,
    damping_lambda=1.0e-3,
)

SR spaces, blocks and AdamSR

SRConfig selects the solve space independently of the numerical backend:

from pynqs.optim import SRConfig, IterativeConfig, TraceConfig

config = SRConfig(
    space="sample",
    solver="iterative",
    damping=1e-3,
    iterative=IterativeConfig(method="cg", tol=1e-5, max_iter=100),
)

lsmr_config = SRConfig(
    space="parameter",
    solver="iterative",
    damping=1e-3,
    iterative=IterativeConfig(method="lsmr", tol=1e-5, max_iter=500, warm_start=True),
)

space is "parameter", "sample" or "auto". Auto selects parameter space when its parameter count is no larger than the global sample count, independently for each block. solver is "direct" or "iterative"; the iterative method is "cg" or "lsmr". LSMR uses parameter space, including when space is auto. LayerSR currently uses direct solves and supports both spaces and auto selection. Different blocks can select different spaces; their aggregate report then uses space="mixed".

A standalone SRConfig defaults to direct/auto, absolute damping 1e-4. Damping must be a positive finite number or a callable returning one. Fixed values are validated at construction, and callable results when solving. A schedule is evaluated once per ordinary SR call or per AdamSR stage, then shared by that call’s blocks; logging does not evaluate it again.

The former SR-specific configuration names are replaced by space, damping, IterativeConfig.fp64_inner_products and AdamSR.vanilla. There are no compatibility aliases or old-checkpoint migration paths.

Ordinary SR is selected by passing sr_config=config to VMCOptimizer with a parameter optimizer such as SGD or Adam. Its low-level entry is pynqs.optim.grad.sr.sr_grad. Penalty SR retains its separate dense workflow. SR requires real trainable parameters, wavefunctions, energies and probabilities. VMCOptimizer rejects store_O_on_cpu=True; CPU Jacobian offloading is currently unsupported by that entry point.

Shared solve path

Ordinary SR and minSR share one implementation. Let A be the centered, probability-weighted log-derivative matrix with shape Ns-by-P, and let b be the centered, weighted energy vector. Their parameter-space and sample-space representations are

\[ \begin{align}\begin{aligned}\text{parameter:}\quad (A^T A+\lambda I)d=A^T b,\\\text{sample:}\quad (AA^T+\lambda I)u=b,\qquad d=A^T u.\end{aligned}\end{align} \]

For positive damping, exact global solves give the same parameter direction. LayerSR applies the corresponding equations separately to each parameter block.

Numerical backends

Backend	Space	Matrix construction
direct	parameter or sample	Forms a P-by-P or Ns-by-Ns Gram matrix and solves it directly.
CG	parameter or sample	Applies the square normal-equation operator without forming a full Gram.
LSMR	parameter	Uses the rectangular Jacobian operator directly, without constructing a square normal-equation operator.

Ordinary and layered SR use a single block loop in sr_grad. AdamSR calls the same solve_sr_block kernel for its vanilla and correction stages. Direct solves pass the Gram matrix, RHS and damping to solve_dense, which solves on rank zero by default. CG uses a normal-equation operator, and LSMR uses the rectangular factor directly. LogDerivative computes the raw log derivatives. Its weighted() method constructs an SRJacobian representing the centered, probability-weighted matrix A; the SR metric is \(S=A^T A\). The object manages explicit storage, matrix-free JVP/VJP actions and caches for that same A. derivatives.py supplies these derivative operations, damping.py resolves the actual shift, and pynqs.linalg performs the numerical solve. Its linear_operator.py contains the shared operators and layouts; cg.py, lsmr.py and direct.py contain their respective solvers. Existing explicit Jacobian storage is bound once for a solve; otherwise the iterative path uses lazy JVP/VJP operations. Ordinary SR and AdamSR share input preparation. Each block’s Jacobian releases its caches on exit, including on failure; switching to explicit storage first releases obsolete VJP graphs.

SRConfig owns the space, damping and partition choices. IterativeConfig only contains the numerical method, tolerance, iteration limit, warm-start switch and inner-product precision. TraceConfig controls trace estimation.

Parameter blocks

layer_sr=True groups trainable parameter tensors by their immediate owning module. For a Transformer this can create separate blocks for projections and normalization modules; it does not automatically group an entire residual block. layer_groups explicitly lists trainable parameter-tensor indices:

config = SRConfig(
    layer_groups=[[0, 2], [1, 3]], space="auto", solver="direct",
    damping=0.03,
)

get_default_layer_groups(model) from pynqs.optim.grad returns the default index groups as lists, ready to inspect, customize and pass to SRConfig.layer_groups.

Groups must be nonempty, disjoint and cover all trainable parameter tensors exactly once. Arbitrary index order and noncontiguous groups are supported. Model-independent checks run when constructing the config; complete coverage is checked when binding the model. VMC logs the selected groups once for ordinary SR or for the shared AdamSR partition.

Each block materializes only its own Jacobian. Sample-space direct solving still needs an Ns-by-Ns Gram matrix; parameter-space solving needs a P_block-by-P_block Gram matrix. Auto avoids a large sample Gram for small blocks. Block diagonalization changes the approximation and does not guarantee a better condition number or an improvement in optimization.

AdamSR

from pynqs.optim import AdamSR, SRConfig, VMCOptimizer

vanilla = SRConfig(
    layer_sr=False,
    space="auto",
    solver="direct",
    damping=1e-3,
)
correction = SRConfig(
    layer_sr=False,
    space="parameter",
    solver="direct",
    damping_mode="relative",
    damping=0.1,
)
opt = AdamSR(
    model.parameters(),
    lr=1.0,
    betas=(0.95, 0.995),
    eps=1e-8,
    max_norm=1.0,
    norm_delay=1000,
    norm_decay_steps=4000,
    vanilla=vanilla,
    correction=correction,
)
# Pass opt to VMCOptimizer; no external sr_config is needed for AdamSR.

AdamSR implements the two-solve update of Algorithm S1 in arXiv:2610.30321. Its defaults are global iterative/sample solves, vanilla absolute damping 1e-3, correction relative damping 0.1 and iterative tolerance 1e-2. Both configs are copied independently. Their layer_sr and layer_groups settings must match exactly, including group order and parameter order within each group. AdamSR checks this at construction and before every solve. Numerical settings, including solve space and damping, remain independent between stages. Configure trainable parameters before constructing the optimizer, and rebuild it if the trainable parameter layout changes.

CI coefficient pre-training uses a separate ordinary optimizer. Construct AdamSR after pre-training; VMCOptimizer.pre_train() rejects an AdamSR optimizer because CI training uses backward/step without an SR solve.

Write A for the centered, probability-weighted log-derivative matrix and b for the centered, weighted energy vector. The first solve produces g. Moments are formed before any clipping:

\[ \begin{align}\begin{aligned}m_t=\mu m_{t-1}+(1-\mu)g_t,\qquad v_t=\beta v_{t-1}+(1-\beta)g_t^2,\\\hat m_t=\frac{m_t}{1-\mu^t},\qquad \hat v_t=\frac{v_t}{1-\beta^t},\qquad D_t=\operatorname{diag}\big[(\hat v_t^{1/4}+\epsilon)^{-1}\big].\end{aligned}\end{align} \]

The second solve uses B=AD and r=b-A mhat, then returns d=mhat+D z:

\[(B^TB+\lambda I)z=B^Tr, \qquad\text{or}\qquad (BB^T+\lambda I)u=r,\quad z=B^Tu.\]

For global exact solves, the resulting direction is equivalently

\[d^\star=\arg\min_d\left\{ \|Ad-b\|^2+\lambda\|D^{-1}(d-\hat m)\|^2 \right\}.\]

The lambda in this expression is the actual correction-stage damping. It regularizes deviations from the moment direction. With negligible epsilon, D inverse squared is diag(sqrt(vhat)); changing the fourth-root exponent changes this metric. The correction depends on b-A mhat. A small residual of the current vanilla solution does not imply a small moment residual, because mhat also contains historical directions.

For LayerSR, each stage uses its parameter-block diagonal approximation. A correction block uses only b-A_block mhat_block. Both stages share the same parameter partition. The implementation completes vanilla, block-local moments and correction consecutively, then releases that block’s Jacobian. L blocks therefore require L explicit Jacobian constructions. Both linear solves remain necessary. Mixed full/block stages and different groupings are rejected.

Damping and precision

Absolute damping uses the configured value directly. Relative damping uses

\[\lambda=c\,\frac{\operatorname{tr}(BB^T)}{N_s},\]

including in parameter space. Each block uses its own trace with the same Ns normalization. Direct solves obtain the current trace from the Gram matrix. Iterative solves use TraceConfig: Hutchinson defaults to 8 probes in batches of 4, seed 0, refreshed every 10 updates. Exact trace estimation is available. Cached trace scales are geometrically smoothed at refresh when the interval exceeds one. Stages and blocks keep independent trace state.

With a current, consistently scaled trace, relative damping makes the exact update invariant under a uniform rescaling of D. A cached trace from an unscaled history need not preserve that identity. Absolute damping intentionally depends on D’s absolute scale.

Jacobian storage and Gram accumulation use dtype_config.real_dtype. Differentiation retains the model precision. Direct factorization and solving use FP64; casting an already accumulated FP32 Gram does not recover FP64 accumulation accuracy. Use the appropriate model and real dtype when that accuracy is required. fp64_inner_products controls iterative scalar inner products only. Direct failures propagate as native PyTorch exceptions. Rank-zero solves share the successful solution, dtype and residual diagnostics with the other ranks.

CG reports the relative residual of its selected square system. LSMR uses SciPy-style stopping tests through the single tol: the solve stops when either test2 = ||B^T r|| / (||B|| * ||r||) <= tol (SciPy’s least-squares test, driven there by atol) or test1 = ||r|| / ||b|| <= tol + tol * ||B|| * ||x|| / ||b|| (SciPy’s compatible-system test, driven by btol). Here ||r|| and ||B|| are SciPy’s running scalar estimates, not fresh operator applications.

Consequently tol is neither SciPy’s atol alone nor atol times btol. The reported rel_res is test2, not an initial-residual reduction ratio. Since test1 can stop the solve first, converged may be true while rel_res exceeds tol.

These norms refer to the problem passed to the iteration core. For a damped warm start this is an augmented correction problem: ||b|| is its shifted RHS norm, and ||x|| is the correction norm. Finite iterative directions are returned with convergence metadata and warnings on non-convergence.

Returned residual_kind is "linear" for direct/CG, "normal_relative" for LSMR, and "block_max" when a report aggregates multiple parameter blocks.

Warm starts and committed state

warm_start=True permits reuse only for iterative parameter-space solves. Sample-space solves always start cold and never gather or save dual initial guesses. Direct solves do not use warm starts. Loading a checkpoint clears all runtime initial guesses; finite-tolerance restarted and uninterrupted runs can therefore differ slightly.

For AdamSR, per-parameter optimizer state contains only first and second moments. Its checkpoint has a separate sr_state containing one update_step and trace metadata keyed by stage and parameter block. Runtime guesses, memory log flags and pending proposals are not serialized. solve() prepares one PendingStep, step() commits it, and zero_grad() discards an uncommitted proposal. reset_runtime_state() clears runtime caches without changing committed moments or trace estimates. Model-only checkpoint loading uses that same reset interface.

AdamSR’s existing norm rule remains lr * min(1, bound / norm(direction)), with optional bound=max_norm / (1 + max(t-norm_delay, 0) / norm_decay_steps) and one-based t. Clipping is applied to the assembled direction after all blocks have completed.

Detailed parameter-update diagnostics in VMCOptimizer are opt-in via optimizer.log_parameter_updates = True. They capture full parameter vectors on rank zero; ordinary updates avoid those snapshots unless DIIS needs them.

TwoStep minSR with LSMR

EnergyMinSRParameterUpdate and FidelityMinSRParameterUpdate accept solver="lsmr". It reuses the matrix-free derivative and damping preparation, then solves the rectangular regularized least-squares problem in parameter space. The existing solver="iterative" continues to select CG.

from pynqs.optim import FidelityMinSRParameterUpdate

parameter_update = FidelityMinSRParameterUpdate(
    solver="lsmr",
    iterative_solve_space="parameter",
    cg_tol=1e-5,
    cg_max_iter=500,
    solver_restart=False,
)

The existing cg_tol, cg_max_iter and solver_restart names also control LSMR. Restart false reuses the previous parameter-space solve vector. CG and LSMR use different stopping metrics, so equal tolerances do not imply equal residuals or direction errors. LSMR currently requires parameter space and cg_preconditioner="none"; MARCH scaling and history shifts are unsupported. Direct and onfly paths retain their existing behavior.

CPU regression tests

From the repository root, or using the runner’s absolute path:

python tests/optim/run_sr_tests.py --ranks 1 2

The runner sets PYTHONPATH, MAX_SORB=64, local Gloo communication and explicit localhost torchrun rendezvous settings. It runs independent SVD comparisons, FP32/FP64 cases, empty ranks, checkpoint transactions, parameter-only warm starts, block fusion and native direct-solve errors. The CPU CI workflow runs the same suite after restoring or building its native extension.

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:

from pynqs.optim import RGNConfig

vmc_opt = VMCOptimizer(
    nqs=model,
    opt=opt,
    sampler_param=sampler_param,
    electron_info=electron_info,
    rgn_config=RGNConfig(
        epsilon=1.0,
        delta=0.0,
        damping_lambda=1.0e-3,
    ),
)

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.

• 10. 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 函数。