netket.experimental.QSR#
- class netket.experimental.QSR[source]#
Bases:
AbstractOptimizationDriverQuantum state reconstruction driver minimizing KL divergence.
This driver variationally reconstructs a target state given the measurement data. It’s achieved by minimizing the average negative log-likelihood, or equivalently, the KL divergence between the distributions given by the data and the variational state:
\[\begin{split}&\min_\theta \frac{1}{N_b} \sum_{b=1}^{N_b} \sum_{\sigma_b} q_b(\sigma_b) \log \left[ \frac{q_b(\sigma_b)}{p_{b\theta}(\sigma_b)} \right] \\ &\approx \min_\theta \frac{1}{N_b} \sum_{b=1}^{N_b} \frac{1}{|D_b|} \sum_{\sigma_b \in D_b} [-\log p_{b\theta}(\sigma_b)],\end{split}\]where \(\theta\) is the variational parameter, \(N_b\) is the number of measurement basis, \(q_b(\sigma_b)\) is the probability of obtaining the outcome state \(\sigma_b\) in the measurement basis \(b\) given the target state, \(p_{b\theta}(\sigma_b)\) is the probability of obtaining the outcome state \(\sigma_b\) in the measurement basis \(b\) given the variational state, and \(D_b\) is the size of the dataset in the measurement basis \(b\).
In practice, the noise introduced by mini-batch training hurts the convergence of accurate quantum state reconstruction. To alleviate this problem, we use a control variate method called stochastic variance reduced gradient (SVRG) to reduce the variance of the gradient estimator. Specifically, we update the parameters \(\theta\) according to
\[\theta_{i+1} = \theta_{i} -\eta \left\{ \underbrace{\nabla_\theta \left[\frac{1}{|B_i|}\sum_{\sigma_b\in B_i} \log p_{b\theta_i}(\sigma_b)\right]}_{\text{I: batch gradient}} - \underbrace{\nabla_\theta \left[\frac{1}{|B_i|}\sum_{\sigma_b\in B_i} \log p_{b\tilde{\theta}_i}(\sigma_b)\right]}_{\text{II: control variate}} + \underbrace{\nabla_\theta \left[\frac{1}{N_b} \sum_{b=1}^{N_b} \frac{1}{|D_b|} \sum_{\sigma_b \in D_b} \log p_{b\tilde{\theta}_i}(\sigma_b)\right]}_{\text{III: expectation of control variate}} \right\},\]where term I is the normal batch gradient, term II is the control variate which is the batch gradient evaluated with a set of previous parameters
\[\begin{split}\tilde{\theta}_i = \begin{cases} \theta_i, &i=0 \mod m, \\ \tilde{\theta}_{i-1}, &\text{otherwise}, \end{cases}\end{split}\]updated for every \(m\) iterations, and term III is the expectation value of the control variate since the mini-batch is sampled uniformly from the whole dataset.
- Inheritance
- __init__(training_data, training_batch_size, optimizer, *, variational_state, preconditioner=<netket.optimizer.preconditioner.IdentityPreconditioner object>, seed=None, batch_sample_replace=True, control_variate_update_freq=None, chunk_size=None)[source]#
Initializes the QSR driver class.
- Parameters:
training_data (
RawQuantumDataset|tuple[list,list]) – A tuple of two arrays (sigma_s, Us). sigma_s is a the sampled states and Us is the corresponding rotations.training_batch_size (
int) – The training batch size.optimizer – The optimizer to use. You can use optax optimizers or choose from the predefined optimizers netket offers.
variational_state (
VariationalState) – The variational state to optimize.preconditioner (
Callable[[VariationalState,Any,Any|None],Any]) – The preconditioner to use. Defaults to identity_preconditioner.batch_sample_replace (
bool|None) – Whether to sample with replacement. Defaults to True.control_variate_update_freq (
None|int|str) – The frequency of updating the control variates. Defaults to None. “Adaptive” for adaptive update frequency, i.e. n_samples // batch size.chunk_size (
int|None) – The chunk size for the control variates. Defaults to None.
- Raises:
- Attributes
- dataset#
- optimizer#
The optimizer used to update the parameters at every iteration.
- preconditioner#
The preconditioner used to modify the gradient.
- state#
Returns the machine that is optimized by this driver.
- step_count#
Returns a monotonic integer labelling all the steps performed by this driver. This can be used, for example, to identify the line in a log file.
- Methods
- KL(target_state=None, n_shots=None)[source]#
Compute average KL divergence loss over a batch of data.
- Parameters:
Warning
Exponentially expensive in the hilbert space size!
Need to know the target state!
- KL_exact(target_state=None, n_shots=1)[source]#
Compute the average KL divergence loss between the variational state and the target state.
- Parameters:
- Return type:
Warning
Exponentially expensive in the hilbert space size!
Need to know the target state!
- KL_whole_training_set(target_state=None, n_shots=None)[source]#
Compute average KL divergence loss over the whole training set.
- Parameters:
Warning
Exponentially expensive in the hilbert space size!
Need to know the target state!
- compute_loss_and_update()[source]#
Performs a step of the optimization driver, returning the PyTree of the gradients that will be optimized.
Concrete drivers must override this method.
Note
When implementing this function on a subclass, you must return the gradient which must match the pytree structure of the parameters of the variational state.
The gradient will then be passed on to the optimizer in order to update the parameters.
Moreover, if you are minimising a loss function you must set the field self._loss_stats with the current value of the loss function.
This will be logged to any logger during optimisation.
- Returns:
the update for the weights.
- entropy(target_state, n_shots=1, no_cache=False)[source]#
Compute the average entropy of the probability distributions given by the target state in different measurement basis.
- Parameters:
- Return type:
Warning
Exponentially expensive in the hilbert space size!
Need to know the target state!
- estimate(observables, fullsum=False)[source]#
Return MCMC statistics for the expectation value of observables in the current state of the driver.
- Parameters:
observables – A pytree of operators for which statistics should be computed.
fullsum (bool)
- Returns:
A pytree of the same structure as the input, containing MCMC statistics for the corresponding operators as leaves.
- nll(return_stats=True)[source]#
Compute the Negative-Log-Likelihood over a batch of data.
Warning
Exponentially expensive in the hilbert space size!
- nll_whole_training_set(return_stats=True)[source]#
Compute the Negative-Log-Likelihood over the whole training set.
Warning
Exponentially expensive in the hilbert space size!
- replace(**kwargs)[source]#
Replace the values of the fields of the object with the values of the keyword arguments. If the object is a dataclass, dataclasses.replace will be used. Otherwise, a new object will be created with the same type as the original object.
- reset()[source]#
Deprecated since version 3.22: Use
reset_step()to reset the sampler state at the beginning of a step. Note that the oldreset()also resetstep_countto 0; this behaviour is no longer supported.
- reset_step(hard=False)[source]#
Resets the state of the driver at the beginning of a new step.
This method is called at the beginning of every step in the optimization.
- Parameters:
hard (
bool) – If True, the reset is a hard reset, resulting in a complete resampling even if resample_fractionNone. (is not)
- run(n_iter, out=(), obs=None, step_size=1, show_progress=True, save_params_every=50, write_every=50, callback=None, timeit=False, _graceful_keyboard_interrupt=True)[source]#
Runs this variational driver, updating the weights of the network stored in this driver for n_iter steps and dumping values of the observables obs in the output logger.
It is possible to control more specifically what quantities are logged, when to stop the optimisation, or to execute arbitrary code at every step by specifying one or more callbacks, which are passed as a list of functions to the keyword argument callback.
Callbacks are functions that follow this signature:
def callback(step, log_data, driver) -> bool: ... return True/False
If a callback returns True, the optimisation continues, otherwise it is stopped. The log_data is a dictionary that can be modified in-place to change what is logged at every step. For example, this can be used to log additional quantities such as the acceptance rate of a sampler.
Alternatively,
AbstractCallbacksubclasses can be used to hook into more stages of the loop. To stop the optimisation early from any callback hook, raiseStopRun: the driver will catch it, finalise all callbacks via theiron_run_endmethod, and return normally without propagating the exception.Loggers are specified as an iterable passed to the keyword argument out. If only a string is specified, this will create by default a
nk.logging.JsonLog. To know about the output format check its documentation. The logger object is also returned at the end of this function so that you can inspect the results without reading the json output.- Parameters:
n_iter (
int) – the total number of iterations to be performed during this run.out (
Iterable[AbstractLog] |None) – A logger object, or an iterable of loggers, to be used to store simulation log and data. If this argument is a string, it will be used as output prefix for the standard JSON logger.obs (
dict[str,AbstractObservable] |None) – An iterable containing all observables that should be computedstep_size (
int) – Every how many steps should observables be logged to disk (default=1)callback (
Union[Callable[[int,dict,AbstractDriver],bool],AbstractCallback,None]) – Callable or list of callable callback functions to stop training given a conditionshow_progress (
bool) – If true displays a progress bar (default=True)save_params_every (
int) – Every how many steps the parameters of the network should be serialized to disk (ignored if logger is provided)write_every (
int) – Every how many steps the json data should be flushed to disk (ignored if logger is provided)timeit (
bool) – If True, provide timing information._graceful_keyboard_interrupt (
bool) – (Internal flag, defaults to True) If True, the driver will gracefully handle a KeyboardInterrupt, usually arising from doing ctrl-C, returning the current state of the simulation. If False, the KeyboardInterrupt will be raised as usual. This only has an effect when running in interactive mode.