tn_noise_learning.trainer

Stochastic gradient descent on the Kraus tensors of an LPDO with optax and JAX.

Trainer

Trainer(config: TrainingConfig | None = None)

Fit an LPDO to a NoiseDataset.

Source code in src/tn_noise_learning/trainer.py

def __init__(self, config: TrainingConfig | None = None) -> None:
    self.config = config or TrainingConfig()

fit

fit(
    model: LPDO,
    train: NoiseDataset,
    test: NoiseDataset | None = None,
    ideal_channel: LPDO | None = None,
) -> History

Train model in place and return the training history.

PARAMETER	DESCRIPTION
model	initial LPDO; on return it holds the parameters of the best epoch. TYPE: LPDO
train	training shots; its circuit is applied inside the loss. TYPE: NoiseDataset
test	held-out shots; split off train when omitted. TYPE: NoiseDataset | None DEFAULT: None
ideal_channel	ground truth to report distance_to_ideal; defaults to train.ideal_channel. TYPE: LPDO | None DEFAULT: None

Source code in src/tn_noise_learning/trainer.py

def fit(
    self,
    model: LPDO,
    train: NoiseDataset,
    test: NoiseDataset | None = None,
    ideal_channel: LPDO | None = None,
) -> History:
    """Train ``model`` in place and return the training history.

    Args:
        model: initial LPDO; on return it holds the parameters of the best epoch.
        train: training shots; its circuit is applied inside the loss.
        test: held-out shots; split off ``train`` when omitted.
        ideal_channel: ground truth to report ``distance_to_ideal``; defaults to ``train.ideal_channel``.
    """
    cfg = self.config
    if test is None:
        train, test = train.split(cfg.test_fraction, cfg.seed)
    ideal = train.ideal_channel if ideal_channel is None else ideal_channel
    if cfg.tp_warmstart_steps > 0:
        model.tensors = project_trace_preserving(model, cfg.tp_warmstart_steps).tensors

    params, skeleton = qtn.pack(model.kraus_network())
    rebuild = lambda p: LPDO.from_network(model.topology, qtn.unpack(p, skeleton))  # noqa: E731
    term_names, weights_at, step, evaluate_loss, opt_state = self._compile(model, params, skeleton, train)
    history = History(num_sites=model.num_sites)
    selection = _Selection(cfg.patience, cfg.tolerance, params)
    global_step = 0
    epoch = 0
    try:
        for epoch in range(1, cfg.epochs + 1):
            sums = np.zeros(len(term_names) + 1)
            for batch in train.batches(cfg.batch_size, shuffle=True, seed=cfg.seed + epoch):
                params, opt_state, total, values = step(
                    params, opt_state, batch.states, batch.povms, weights_at(global_step)
                )
                sums += np.concatenate([[float(total)], np.asarray(values)])
                global_step += 1
                if cfg.tp_project_every and global_step % cfg.tp_project_every == 0:
                    projected = project_trace_preserving(rebuild(params), cfg.tp_project_steps)
                    params = dict(zip(params, qtn.pack(projected.kraus_network())[0].values(), strict=True))
            num_batches = len(train) // cfg.batch_size
            test_loss = np.mean(
                [
                    float(evaluate_loss(params, b.states, b.povms, weights_at(global_step))[0])
                    for b in test.batches(min(cfg.batch_size, len(test)), shuffle=False, drop_last=False)
                ]
            )
            metrics = {"train_loss": sums[0] / num_batches, "test_loss": float(test_loss)}
            metrics |= {name: sums[i + 1] / num_batches for i, name in enumerate(term_names)}
            current = rebuild(params)
            metrics |= self._model_metrics(current, ideal)
            history.append(metrics, current.tensors if cfg.snapshot_params else None)
            if epoch % cfg.log_every == 0:
                logger.info(
                    "epoch %3d/%d  %s", epoch, cfg.epochs, "  ".join(f"{k}={v:.4g}" for k, v in metrics.items())
                )
            if cfg.monitor not in metrics:
                raise KeyError(f"monitor {cfg.monitor!r} is not a recorded metric: {sorted(metrics)}")
            if selection.update(metrics[cfg.monitor], params, epoch):
                break
    except KeyboardInterrupt:
        logger.warning("Training interrupted at epoch %d; keeping the best model so far.", epoch)

    model.tensors = rebuild(selection.best_params).tensors
    history.best_epoch = selection.best_epoch
    return history

TrainingConfig dataclass

TrainingConfig(
    epochs: int = 50,
    batch_size: int = 256,
    learning_rate: ScalarOrSchedule = (
        lambda: optax.exponential_decay(
            0.01, 200, 0.8, end_value=0.0001
        )
    )(),
    optimizer: Callable[
        [ScalarOrSchedule], GradientTransformation
    ] = optax.adam,
    tp_weight: Schedule = 100.0,
    tp_warmstart_steps: int = 500,
    tp_project_every: int | None = None,
    tp_project_steps: int = 50,
    extra_loss_terms: dict[
        str, tuple[LossTerm, Schedule]
    ] = dict(),
    metrics: dict[str, Metric] = dict(),
    monitor: str = "test_loss",
    patience: int = 10,
    tolerance: float = 0.0001,
    test_fraction: float = 0.1,
    seed: int = 0,
    snapshot_params: bool = False,
    log_every: int = 1,
    contract_optimize: str = "auto-hq",
)

Hyper-parameters of Trainer.

PARAMETER	DESCRIPTION
epochs	maximum number of passes over the training set. TYPE: int DEFAULT: 50
batch_size	shots per gradient step; every batch has exactly this size. TYPE: int DEFAULT: 256
learning_rate	constant, optax schedule, or callable of the step index; defaults to an exponential decay from 1e-2 by a factor 0.8 every 200 steps down to 1e-4. TYPE: ScalarOrSchedule DEFAULT: (lambda: exponential_decay(0.01, 200, 0.8, end_value=0.0001))()
optimizer	factory turning the learning rate into an optax optimizer. TYPE: Callable[[ScalarOrSchedule], GradientTransformation] DEFAULT: adam
tp_weight	weight of the trace-preservation penalty, constant or callable of the step index. The trained trace exceeds 2**n by about 1/(2 tp_weight) relative; weights far above 1000 make Adam stiff. TYPE: Schedule DEFAULT: 100.0
tp_warmstart_steps	L-BFGS-B steps moving the initial model onto the trace-preserving manifold (0 disables). TYPE: int DEFAULT: 500
tp_project_every	re-run the L-BFGS-B projection every this many steps (None disables). TYPE: int | None DEFAULT: None
tp_project_steps	L-BFGS-B steps per projection. TYPE: int DEFAULT: 50
extra_loss_terms	name -> (term, weight) added to the loss; a term maps (lpdo_tn, probs) to a scalar. TYPE: dict[str, tuple[LossTerm, Schedule]] DEFAULT: dict()
metrics	name -> metric(model) evaluated on the current model after every epoch. TYPE: dict[str, Metric] DEFAULT: dict()
monitor	metric used for early stopping and model selection. TYPE: str DEFAULT: 'test_loss'
patience	epochs without improvement of monitor before training stops. TYPE: int DEFAULT: 10
tolerance	minimum decrease of monitor that counts as an improvement. TYPE: float DEFAULT: 0.0001
test_fraction	held-out fraction used when fit receives no test set. TYPE: float DEFAULT: 0.1
seed	seed for the train/test split and batch shuffling. TYPE: int DEFAULT: 0
snapshot_params	record the Kraus tensors after every epoch in the history. TYPE: bool DEFAULT: False
log_every	log one line every this many epochs. TYPE: int DEFAULT: 1
contract_optimize	cotengra path optimizer for the loss contraction. TYPE: str DEFAULT: 'auto-hq'

evaluate_schedule

evaluate_schedule(schedule: Schedule, step: int) -> float

Value of a constant or step-dependent schedule.

Source code in src/tn_noise_learning/trainer.py

def evaluate_schedule(schedule: Schedule, step: int) -> float:
    """Value of a constant or step-dependent schedule."""
    if isinstance(schedule, Callable):
        return float(schedule(step))
    return float(schedule)

project_trace_preserving

project_trace_preserving(model: LPDO, steps: int) -> LPDO

Minimise the trace-preservation penalty alone with L-BFGS-B, starting from model.

Source code in src/tn_noise_learning/trainer.py

def project_trace_preserving(model: LPDO, steps: int) -> LPDO:
    """Minimise the trace-preservation penalty alone with L-BFGS-B, starting from ``model``."""
    topology = model.topology
    optimizer = qtn.TNOptimizer(
        model.kraus_network(),
        loss_fn=lambda tn: tp_penalty(full_network(tn, topology)),
        optimizer="L-BFGS-B",
        autodiff_backend="jax",
        jit_fn=True,
        progbar=False,
    )
    return LPDO.from_network(topology, optimizer.optimize(steps))