Test-Time Training with Quantum Auto-Encoder: From Distribution Shift to Noisy Quantum Circuits

Damien Jian^*, Yu-Chao Huang^*, Hsi-Sheng Goan

National Taiwan University
^*Equal contribution. Authors are listed in reverse alphabetical order.

Test-Time Training with Quantum Auto-Encoder (QTTT) as a Flexible Plug-and-Play Extension for Quantum Neural Networks (QNNs). QTTT adapts to distribution shift and mitigates random unitary noise present during inference by minimizing the self-supervised loss from the quantum auto-encoder. QTTT integrates easily with existing QNNs, with a mild computational overhead ratio during training and test-time training that is inversely proportional to the main QNN depth. QTTT serves as a bridge toward the practical utility of QNNs in the future.

Abstract

In this paper, we propose test-time training with the quantum auto-encoder (QTTT). QTTT adapts to data distribution shifts between training and testing data and quantum circuit error by minimizing the self-supervised loss of the quantum auto-encoder. Empirically, we show that QTTT is robust against data distribution shifts and effective in mitigating random unitary noise in the quantum circuits during the inference. Additionally, we establish the theoretical performance guarantee of the QTTT architecture. Our novel framework presents a significant advancement in developing quantum neural networks for future real-world applications and functions as a plug-and-play extension for quantum machine learning models.

Methodology

QTTT features a Y-shaped QNN architecture. The main branch is a data re-uploading classifier for classification tasks, with the corresponding loss function \( \ell_{\text{ML}} (\cdot; \theta_{\text{L}}, \theta_{\text{E}}, \theta_{\text{M}}) \). The auxiliary branch is a QAE, with the corresponding loss function \( \ell_{\text{AE}} (\cdot; \theta_{\text{L}}, \theta_{\text{E}}, \theta_{\text{D}}) \). We use a multi-task objective during training, with trainable parameters \( (\theta_{\text{L}}, \theta_{\text{E}}, \theta_{\text{D}}, \theta_{\text{M}}) \). In the test-time training stage, the shared parameters \( (\theta_{\text{L}}, \theta_{\text{E}}) \) in the linear layer and encoder enable QTTT to capture distribution shifts in test data or quantum circuit errors.

The Overall Architecture of Test-Time Training with Quantum Auto-Encoder (QTTT).

The general rotation layer (left) consists of a series of parameterized general single qubit gates \( U_3 \) with three trainable parameters. The entanglement layer (right) consists of a series of ladder CNOT gates for strong entanglement.

The QAE (top) consists of an amplitude encoding block, an encoder, a multi-qubit SWAP operation, and a decoder. We make amplitude measurements at the end. The main task branch (bottom) consists of the same amplitude encoding and the encoder in the QAE, followed by a data re-uploading block. We make Pauli-Z measurements at the end. The SWAP operations discard trash qubits by swapping the last two registers (both with Nt qubits) of the quantum circuits.

Experiments

Test Accuracy \((\%)\) on QML-Benchmark with Different Levels of Noise (\( d_x = 5\)).

Test Accuracy \((\%)\) on QML-Benchmark with Different Levels of Noise (\( d_x = 10\)).

Test Accuracy Improvements \((\%)\) for Noisy Quantum Circuits.

We select 12 noise level upper bounds \( \epsilon_\theta \) from \( \pi/40\) to \( 12\pi/40 \), applying random rotation gates with angle sampled from \( \text{Unif}(0,\epsilon_\theta) \) after each layer of the QML models. QTTT-Online and QTTT-Batch improve the model accuracy by \(7.0\%\) and \(5.9\%\), respectively, compared to the vanilla data re-uploading model.

Ablation Studies of QTTT. The best and the second best accuracy across different corruption types and quantum circuit noise level upper bounds are shown in bold and with underlining, respectively.

BibTeX

@misc{jian2024qttt,
  title={Test-Time Training with Quantum Auto-Encoder: From Distribution Shift to Noisy Quantum Circuits}, 
  author={Damien Jian and Yu-Chao Huang and Hsi-Sheng Goan},
  year={2024},
  eprint={2411.06828},
  archivePrefix={arXiv},
  primaryClass={quant-ph},
  url={https://arxiv.org/abs/2411.06828}, 
}