# Dr.ICL: Demonstration-Retrieved In-context Learning

Man Luo1 Xin Xu2 Zhuyun Dai2 Panupong Pasupat2
Mehran Kazemi2 Chitta Baral1 Vaiva Imbrasaite2 Vincent Y Zhao2
1 Arizona State University 2 Google Research
{mluo26<EMAIL_ADDRESS>
{xxujasmine, zhuyundai, ppasupat, mehrankazemi, vimbrasaite<EMAIL_ADDRESS>
###### Abstract

In-context learning (ICL), teaching a large language model (LLM) to perform a
task with few-shot demonstrations rather than adjusting the model parameters,
has emerged as a strong paradigm for using LLMs. While early studies primarily
used a fixed or random set of demonstrations for all test queries, recent
research suggests that retrieving semantically similar demonstrations to the
input from a pool of available demonstrations results in better performance.
This work expands the applicability of retrieval-based ICL approaches by
demonstrating that even simple word-overlap similarity measures such as BM25
outperform randomly selected demonstrations. Furthermore, we extend the
success of retrieval-based ICL to instruction-finetuned LLMs as well as Chain-
of-Thought (CoT) prompting. For instruction-finetuned LLMs, we find that
although a model has already seen the training data at training time,
retrieving demonstrations from the training data at test time yields better
results compared to using no demonstrations or random demonstrations. Last but
not least, we train a task-specific demonstration retriever that outperforms
off-the-shelf retrievers.

Dr.ICL: Demonstration-Retrieved In-context Learning

Man Luo1 Xin Xu2 Zhuyun Dai2 Panupong Pasupat2 Mehran Kazemi2 Chitta Baral1
Vaiva Imbrasaite2 Vincent Y Zhao2 1 Arizona State University 2 Google Research
{mluo26<EMAIL_ADDRESS>{xxujasmine, zhuyundai, ppasupat, mehrankazemi,
vimbrasaite<EMAIL_ADDRESS>

##  1 Introduction

Language models are now the foundation models for many natural language
processing tasks across a wide range of domains (Bommasani et al., 2021). One
of the most exciting emergent abilities (Wei et al., 2022a) of large language
models (LLMs) is in-context learning (ICL) (Brown et al., 2020; Mishra et al.,
2022). With ICL, instructions and a few demonstrative examples are augmented
to the inputs to LLMs, allowing them to perform well on new tasks without the
need for fine-tuning.

Figure 1: The average performance of PaLM and Flan-PaLM on five datasets, with
one and few-shot ICL. Retrieved demonstrations given by either BM25 or GTR
yield better performance than random demonstrations.

Typically, ICL approaches utilize random or hand-crafted demonstrations that
are applied across various queries. This may, however, not always be optimal.
Recent research has revealed that using demonstrations semantically similar to
the input query can enhance performance Liu et al. (2022). In this work, we
investigate two off-the-shelf retrievers, BM25 Robertson et al. (2009) and GTR
Ni et al. (2021), where BM25 is a sparse retriever that finds demonstrations
with the highest (weighted) word overlap with the query, while GTR is a dense
retriever that seeks demonstrations semantically closest to the query. Then,
we utilize them to obtain query-specific demonstrations, and study
demonstration-retrieved ICL (Dr. ICL) with a general and an instruction-
finetuned LLM. Beyond previous work, several interesting findings are
discovered through our experiments as shown in Figure 1. Firstly, we establish
that both BM25 and GTR can find more effective demonstrations than random
demonstrations in both one-shot and few-shot ICL settings. Such off-the-shelf
retrievers make Dr. ICL an appealing paradigm for real-world applications.
Secondly, our results with an instruction-finetuned LLM, i.e., Flan-PaLM
(Chung et al., 2022), indicate that training data can be useful not only for
training models but for accompanying a retriever for testing, suggesting a
more efficient way to utilizing training data which are expensive to collect.
Lastly, by combining with an advanced prompting technique, Chain-of-Thought
(CoT) Wang et al. (2022), demonstration-retrieved proves to be more effective
than relying solely on CoT. This suggests that Dr. ICL can boost the
performance of powerful prompt engineering techniques.

Next, we aim to go beyond off-the-shelf retrievers which are often geared
towards question answering or information retrieval tasks thus the retrieved
demonstrations might capture query-specific knowledge required to answer the
query. However, the retrieved demonstrations given by the off-the-shelf
retrievers might not represent the nature of the task and how the task should
be solved in general. Consider, for example, the query “In a barn are chickens
and rabbits with 35 heads and 94 legs total. How many chickens and rabbits are
there?”. Off-the-shelf retrievers may mostly provide information about the
animals in the question and their properties such as number of heads and legs
(i.e. query-specific knowledge), but may not provide enough similar linear
algebra questions (i.e. information about the nature of the task).

Therefore, we develop a demonstration retriever that is tailored to retrieving
representative demonstrations. Figure 2 showcases the process of training the
demonstration retriever: we first create a demonstration retrieval training
set using signals from a language model. Concretely, we use an off-the-shelf
retriever to find demonstration candidates for a given input question, prepend
them to the question, and then obtain probabilities from the language model to
re-rank the candidates. We then use the top-$n$ and bottom-$n$ candidates as
positive and hard-negative examples, respectively, to construct a training set
and train the retriever to identify the best demonstration example for a given
query. Experimental results show that the demonstration retriever outperforms
off-the-shelf retrievers, with more noticeable improvement in one-shot ICL.
This encouraging result indicates that the trained retriever could offer an
effective substitute for off-the-shelf models.

##  2 Related Work

Paper |  LLMs |  Retrieval Method |  Retrieval Corpus |  Evaluation Tasks |  # of Shots in Prompts |  CoT
---|---|---|---|---|---|---
KATE 2022 |  GPT-3 |  RoBERTa+kNN |  In-Domain TD |  SA, T2T |  Few-shots |  No
EPR 2022 |  GPT-J, GPT-Neo, CODEX, GTP-3 |  SBERT, BM25, FT Retriever |  In-Domain TD |  SRM |  Few-shots |  No
CEIL 2023 |  GPT-Neo, GPT2-XL, CodeX |  BM25, BERT, DPR, FT Retriever |  In-Domain TD |  SA, PD, NLI, CSR, QA, codeG, and SP |  Few shots |  No
UPRISE 2023 |  GPT-Neo, BLOOM, OPT, GPT-3 |  FT Retriever |  Cross Tasks TD |  RC, QA, NLI, SA, CSR, CR, PD |  Few shots |  No
Ours |  PaLM, Flan-PaLM |  BM25, GTR, FT Retriever |  In-Domain TD |  QA, NLI, MathR, BC |  One-shot, Few-shots |  Yes

Table 1: Comparison with Related Work. TD: training data, QA: question
answering, RC: reading comprehension, NLI: natural language inference, SA:
sentiment analysis, CSR: commonsense reasoning, CR: Coreference Resolution,
MathR: mathmatical reasoning, PD: paraphrase detection, SP:semantic parsing,
CodeG: code generation, SRM: Sentence representation mapping, T2T: Table to
Text generation, Question Answering,

###  2.1 Few-shot In-context Learning

Few-shot in-context learning (ICL) is a technique that allows language models,
such as GPT-3 Brown et al. (2020) and PaLM Chowdhery et al. (2022), to
generalize to new tasks based on a small number of examples. ICL offers
several advantages over the traditional training approach of language models,
which involves pre-training followed by fine-tuning. One key benefit is that
fine-tuning may not always be feasible due to restricted access to the LLM or
inadequate computational resources (Brown et al., 2020). Additionally, ICL
avoids the issues commonly associated with fine-tuning, such as overfitting or
shocks Ying (2019); Kazemi et al. (2023), as it does not modify the model’s
parameters, allowing it to remain general. However, the effectiveness of ICL
varies depending on various factors, such as the order of the demonstrations
Kumar and Talukdar (2021), the distribution of the demonstrations Min et al.
(2022), and the complexity and quality of the prompts themselves Zhao et al.
(2021); Arora et al. (2022). Some research has shown that lower perplexity
prompts Gonen et al. (2022) and open-ended question-answer formats Arora et
al. (2022) tend to lead to better performance, while others have found that
intermediate reasoning steps Wei et al. (2022b) and higher complexity prompts
Fu et al. (2022) can also improve results on certain tasks Suzgun et al.
(2022); Wang et al. (2022). In an effort to understand how ICT works, studies
have suggested that ICL may involve implicit Bayesian inference Xie et al.
(2021) and a symbiotic relationship between text and patterns Madaan and
Yazdanbakhsh (2022), and can behave similarly to explicit fine-tuning Dai et
al. (2022). Our work focus on the effect of demonstrations for ICL with large
language models.

###  2.2 Retrieval Augmented Demonstrations

As summarized in Table 1, several previous works have explored retrieval
techniques for identifying more informative demonstrations to boost in-context
learning. KATE Liu et al. (2022) discovers that semantically closer
demonstrations outperform random ones for GPT-3 in-context learning. They
employ language models trained on tasks like natural language inference and
sentence textual similarity as semantic representations and utilize the kNN
algorithm to search for demonstrations. EPR Rubin et al. (2022) develops a
retriever based on language model signals to find superior demonstrations
compared to off-the-shelf retrievers. Instead of using a separate retriever
for each task, UPRISE Cheng et al. (2023) merges multiple training datasets
into a retrieval corpus and trains a universal retriever for cross-domain
tasks. PARC Nie et al. (2022) employs a multilingual retrieval strategy to
find demonstrations from high-resource tasks, thereby enhancing the
performance of low-resource domain tasks. CEIL Ye et al. (2023), instead of
retrieving few-shot demonstrations independently, introduces an iterative
retrieval method to identify both diverse and similar few-shot examples. While
the aforementioned methods retrieve demonstrations from training data, Madaan
et al. (2022); Dalvi et al. (2022) incorporate human feedback to create
demonstrations and maintain a dynamic retrieval corpus. Z-ICL Lyu et al.
(2022) generates pseudo demonstrations to enhance zero-shot in-context
performance. In contrast to the methods that retrieve explicit demonstrations,
RETROPROMPT Chen et al. (2022) transforms explicit demonstrations into
implicit neural demonstrations represented by vectors. Rather than using a
retriever, Ram et al. (2023) applies a cross-attention reranker to re-rank
documents retrieved by BM25.

##  3 Demonstration-Retrieved In-Context Learning (Dr. ICL)

Figure 2: Pipeline for training demonstration retriever and inference (R for a
neural retriever). Figure on the left shows the procedure of obtaining data to
train a demonstration retriever: an off-the-shelf retriever takes an input
query $x_{q}$ and retrieves top-$k$ (e.g., 100) demonstrations candidates from
the training corpus. Then an LLM is used to output the score of the ground
truth of $y_{q}$ with each retrieved demonstration and $x_{q}$. Figure on the
right shows the inference pipeline for in-context learning with the trained
demonstration retriever.

We start by describing ICL for general tasks (including classification or
generation tasks). For a task $T$, given an input text $x_{q}$, an LLM is used
to predict the answer $y_{q}$ conditioned on a set of _demonstrations_ of the
task, $Demo=\\{d_{1},d_{2},\cdots,d_{n}\\}$, where $d_{i}=(x_{i},y_{i})$ is a
pair of input and ground truth answer. Typically, $d_{i}$ is linearized as a
string (e.g., “question: $x_{i}$ \n answer: ${y_{i}}$”) and then provided to
the LM. Recently, the Chain-of-thoughts prompting technique Wei et al. (2022b)
has demonstrated its effectiveness in handling complex reasoning tasks. The
primary concept involves including intermediate reasoning steps for each
demonstration, so it consists of not only the input and ground truth answer
but also the step-by-step reasoning process.

There are multiple strategies for choosing the set of demonstrations. For
instance, one could randomly or manually select a fixed set $Demo$ to be
applied to all queries of task $T$. Alternatively, a retriever can be used to
find query-specific demonstrations from the training set $D_{train}$:

$Demo_{x_{q}}=Retriever(x_{q},D_{train},n),$ (1)

where $Demo_{x_{q}}$ are the top-$n$ demonstrations that the retriever
considers most suitable for the input $x_{q}$. In this work, we consider two
off-the-shelf retrievers, BM25 and GTR (Section 3.1), and then propose a
method to train a retriever tailored to the target task $T$ (Section 3.2).

###  3.1 Off-the-shelf Retrievers

BM25 Robertson et al. (2009) is a bag-of-words model that calculates relevance
scores using term frequency, inverse document frequency, and document length
normalization. It has proven effective and efficient, making it easily
deployable in large-scale, real-world applications. However, BM25 heavily
relies on keyword matching and lacks context understanding, which may result
in less accurate outcomes. In contrast, GTR Ni et al. (2021) is a dual-encoder
neural retriever (based on T5) trained on the MS Marco dataset Nguyen et al.
(2016). GTR excels in semantic and context comprehension and is easily
transferable to downstream tasks or specific domains. However, it has lower
memory and computational efficiency, and lacks interpretability.

###  3.2 Demonstration Retriever Training

Demonstration retrieval aims to find the most representative demonstrations
for each input query. Ideally, the demonstrations should capture both (a) the
query-specific knowledge required to answer the query, and (b) the nature of
the task and how the task should be solved in general.

Off-the-shelf retrievers such as BM25 and GTR were designed for information
retrieval and question answering. As such, they mostly retrieve demonstrations
of type (a) but not (b). To fill this gap, we propose to train a demonstration
retriever by leveraging the feedback from a language model. As demonstrated in
Figure 2, the process involves two steps: obtaining the training data and
training a retriever on the data.

#### Obtain the Training data

We want to teach the retriever model to locate examples that lead to the most
accurate predictions. We propose to mine a set of demonstrations for each
input query $x_{q}$ in the training data as follows. First, given a question-
answer pair $(x_{q},y_{q})\in D_{train}$, we use an off-the-shelf retriever to
find a demonstration candidate set $D$ for $x_{q}$, where $x_{q}$ is exclusive
from $D$. Second, we test each demonstration $d\in D$ on how much it helps on
the target task. The LM probability $p_{\text{LM}}(y_{q}\mid d,x_{q})$ of the
gold answer $y_{q}$ is used as the score for the demonstration. Finally, we
keep the top-$n$ demonstration as the positive demonstrations, and the
bottom-$n$ as the hard negative demonstrations.

#### Training Procedure

Our retriever is a dual encoder, which defines the score of any query-document
pair $(q,d)$ as $s(q,d)=v_{q}^{\top}v_{d}$, where $v_{q}$ and $v_{d}$ are the
embeddings of $q$ and $d$. We initialize our retriever with GTR, and then
fine-tune it on the training data via contrastive loss with both in-batch and
hard negatives:

$\mathcal{L}_{con}=-\log\frac{e^{s(q,d^{+})}}{e^{s(q,d^{+})}+\sum_{j}e^{s(q,d_{j}^{-})}},$
(2)

where $d^{+}$ and $d_{j}^{-}$ are the positive and negative demonstrations.
The negative demonstrations include the positive demonstrations for the other
input queries in the same batch and 1 randomly-chosen hard negative
demonstration.

Figure 3: PaLM: One-shot and few-shot inference with three types of
demonstrations, random, BM25, and GTR.  Figure 4: Flan-PaLM: One-shot and few-
shot inference with three types of demonstrations, random, BM25, and GTR.

##  4 Experiments

#### Datasets and Evaluation Metrics

We study various tasks across 5 datasets: free-form question answering (NQ),
natural language inference (ANLI-r3), mathematical reasoning (GSM8k and AQuA)
and boolean question answering (StrategyQA). For the last three datasets, we
apply CoT. All the tasks are evaluated by exact matching accuracy.

#### Language Models

PaLM-540B Chowdhery et al. (2022) and Flan-PaLM (540B) Chung et al. (2022) are
used as the primary LLMs. Both models have the same architecture, but Flan-
PaLM has been further trained on thousands of tasks for instruction learning
(including all the five datasets we studied in this paper) and shows superior
generalization performance compared to PaLM. At inference time, we use the
temperature of 0.0 and maximum decoding length 10 for tasks without CoT and
256 for tasks involving CoT.

#### Retrievers

As explained in §3, we explore using BM25 and GTR as off-the-shelf retrievers,
as well as training our own retriever for each task.

For BM25, we use uncased BERT wordpiece tokenization and parameters
$(k_{1},b)=(1.5,0.75)$. For GTR, we use the pretrained GTR-Base model.

When mining data for training our retriever, we use the pretrained GTR to
retrieve 100 demonstrations candidates, and then use PaLM-62B to score each
candidate. (We used the smaller PaLM-62B instead of 540B for efficiency.) Then
we select the top-5 reranked demonstrations as the positive candidates to
fine-tune GTR.

#### Retrieval Corpus

We create a separate retrieval corpus for each task using the associated
training data. For tasks with CoT, each entry in the corpus is composed of the
question, the CoT, and the answer, while for other tasks are without the CoT.

###  4.1 Results

#### Off-the-shelf-retriever performance

Figures 3 and 4 show the performance of PaLM and Flan-PaLM under one-shot and
few-shot ICL settings, with and without retrievers. We make the following
observations.

Observation 1: Off-the-shelf retrievers are capable of finding more effective
demonstrations than random ones. Figure 3 shows that the demonstrations
retrieved by BM25 or GTR are better than random ones under both one-shot and
few-shot scenarios for the PaLM model. It is worth mentioning that BM25 is
more efficient in terms of indexing memory and retrieval latency compared to
semantic retrievers like GTR or other sentence encoders Liu et al. (2022),
which makes it easier to deploy.

Observation 2: Dr. ICL improves instruction-finetuned LLM. Previous research
has primarily focused on investigating demonstration retrieved ICL with
general LLMs (such as GPT-3) rather than instruction-finetuned LLMs, possibly
because they did not consider reusing the training data. In our study, we
examine Dr. ICL with Flan-PaLM, an instruction-finetuned LLM, and present the
results in Figure 4. Overall, the retrieved demonstrations outperform no
demonstrations or random demonstrations. This implies that the training data
should be reused during inference as they can be retrieved and enhance the
performance, even if the model has already seen such data. We conjecture that
the retrieved demonstrations may enhance knowledge localization for ICL, which
could explain the observed improvement.

Observation 3: Dr. ICL can further improve advanced prompting technique,
Chain-of-Thought. In our experiments on GSM8k, StrategyQA, and AQuA, using Dr.
ICL in conjunction with CoT results in improved performance under both one-
shot and few-shot ICL scenarios. This finding suggests that Dr. ICL has the
potential to enhance the performance of powerful prompting techniques.

The observations above hold significant values for real-world applications.
Incorporating ICL with a simple BM25 demonstration retriever, which is highly
scalable in terms of latency and indexing memory, is proven to improve the
performance of the LLM, including when instruction finetuning or Chain-of-
Thought were used. Examples of retrieved demonstrations given by the off-the-
shelf retrievers are given in the Table 4 in Appendix.

#### Trained Demonstration Retriever Performance

Task | Method | One Shot | Few Shots
---|---|---|---
NQ | GTR | 37.8 | 43.9
Demo-GTR(our) | 39.2(+1.4) | 43.9
ANLI (r3) | GTR | 54.0 | 59.0
Demo-GTR(our) | 54.8(+0.8) | 59.0
GSM8k | GTR | 57.7 | 61.0
Demo-GTR(our) | 59.3(+1.6) | 61.5(+0.5)
Avg. | GTR | 49.8 | 54.6
Demo-GTR(our) | 51.1(+1.3) | 54.8(+0.2)

Table 2:  Performance of PaLM using GTR and Demo-GTR retrieved demonstrations.
Demo-GTR consistently achieves better performance than GTR in one-shot case.

We experiment our trained demonstration retriever with PaLM. Table 2 displays
both one-shot and few-shot performance and show that the demonstration
retriever is better than off-the-shelf GTR in almost all cases, leading to a
better overall performance. Notably, the improvements were most significant in
the one-shot ICL scenario, which requires less inference latency and computing
resources than few-shot ICL. These promising results suggest that the trained
retriever could provide an effective alternative to off-the-shelf models.

##  5 Analysis

To rule out the chance that retrieved demonstrations are more advantageous
than random ones simply because in the benchmark datasets the former’s answers
are identical to the correct ones, we assess the overlap percentage between
the demonstration responses and the target. In the few-shot scenario, we
aggregate the answers from the demonstrations via majority voting. From Table
3, it is evident that for the first forth datasets, the overlap ratio is
roughly equal to or less than the uniform distribution, suggesting that the
benefits of the retrieved demonstrations are not due to label identification.
In the case of the NQ, we notice a considerable overlap between demonstration
answers and the ground truth. We then randomly select 100 instances out of the
433 overlapped cases from GTR-retrieved demonstrations (one-shot) and manually
examine them. We find that, indeed, for the majority of the 100 instances, the
input questions are semantically equal to the demonstration questions.

Task | Random | Retriever | One-shot | Few-shot
---|---|---|---|---
ANLI3 | 33.33 | BM25 | 33.33 | 31.42
GTR | 34.75 | 32.25
StrategyQA | 50.0 | BM25 | 48.79 | 47.34
GTR | 47.83 | 48.31
AQUA | 20.0 | BM25 | 22.83 | 25.98
GTR | 24.02 | 22.05
GSM8K | 0.0 | BM25 | 1.36 | 1.82
GTR | 0.99 | 1.14
NQ | 0.0 | BM25 | 8.95 | 8.70
GTR | 11.99 | 11.08
Table 3: Overlapped Ratio of Demonstrations Answers with Targets: Random
represents the probability of selecting the correct label if we select
randomly from the space of possible labels.

##  6 Discussion and Conclusion

In this work, we first leverage two off-the-shelf retrievers to enhance ICL by
searching query-oriented demonstrations. Our experiments demonstrated that
off-the-shelf retrievers are more effective than random demonstrations, with
GTR generally retrieving more representative demonstrations than BM25. More
importantly, our results with Flan-PaLM indicated that training data can be
useful not only for training a model but also for improving the performance of
fine-tuned LLM during testing via ICL. Our experiments with CoT also suggests
that integrating Dr. ICL with advanced prompting techniques can further
improve model’s performance. Additionally, we trained a demonstration
retriever that further improved the overall performance of off-the-shelf
retrievers, with the most significant improvements observed in the one-shot
scenario. One interesting future research challenge is retrieving
demonstrations across tasks in situations where training data is not
available.

## References

  * Arora et al. (2022) Simran Arora, Avanika Narayan, Mayee F Chen, Laurel J Orr, Neel Guha, Kush Bhatia, Ines Chami, Frederic Sala, and Christopher Ré. 2022.  Ask me anything: A simple strategy for prompting language models.  _arXiv preprint arXiv:2210.02441_.
  * Bommasani et al. (2021) Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. 2021.  On the opportunities and risks of foundation models.  _arXiv preprint arXiv:2108.07258_.
  * Brown et al. (2020) Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901.
  * Chen et al. (2022) Xiang Chen, Lei Li, Ningyu Zhang, Xiaozhuan Liang, Shumin Deng, Chuanqi Tan, Fei Huang, Luo Si, and Huajun Chen. 2022.  Decoupling knowledge from memorization: Retrieval-augmented prompt learning.  In _Advances in Neural Information Processing Systems_.
  * Cheng et al. (2023) Daixuan Cheng, Shaohan Huang, Junyu Bi, Yuefeng Zhan, Jianfeng Liu, Yujing Wang, Hao Sun, Furu Wei, Denvy Deng, and Qi Zhang. 2023.  Uprise: Universal prompt retrieval for improving zero-shot evaluation.  _arXiv preprint arXiv:2303.08518_.
  * Chowdhery et al. (2022) Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. 2022.  Palm: Scaling language modeling with pathways.  _arXiv preprint arXiv:2204.02311_.
  * Chung et al. (2022) Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Eric Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, et al. 2022.  Scaling instruction-finetuned language models.  _arXiv preprint arXiv:2210.11416_.
  * Dai et al. (2022) Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Zhifang Sui, and Furu Wei. 2022.  Why can gpt learn in-context? language models secretly perform gradient descent as meta optimizers.  _arXiv preprint arXiv:2212.10559_.
  * Dalvi et al. (2022) Bhavana Dalvi, Oyvind Tafjord, and Peter Clark. 2022.  Towards teachable reasoning systems.  _arXiv preprint arXiv:2204.13074_.
  * Fu et al. (2022) Yao Fu, Hao Peng, Ashish Sabharwal, Peter Clark, and Tushar Khot. 2022.  Complexity-based prompting for multi-step reasoning.  _arXiv preprint arXiv:2210.00720_.
  * Gonen et al. (2022) Hila Gonen, Srini Iyer, Terra Blevins, Noah A Smith, and Luke Zettlemoyer. 2022\.  Demystifying prompts in language models via perplexity estimation.  _arXiv preprint arXiv:2212.04037_.
  * Kazemi et al. (2023) Mehran Kazemi, Sid Mittal, and Deepak Ramachandran. 2023.  Understanding finetuning for factual knowledge extraction from language models.  _arXiv preprint arXiv:2301.11293_.
  * Kumar and Talukdar (2021) Sawan Kumar and Partha Talukdar. 2021.  Reordering examples helps during priming-based few-shot learning.  In _Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021_ , pages 4507–4518, Online. Association for Computational Linguistics.
  * Liu et al. (2022) Jiachang Liu, Dinghan Shen, Yizhe Zhang, William B Dolan, Lawrence Carin, and Weizhu Chen. 2022.  What makes good in-context examples for gpt-3?  In _Proceedings of Deep Learning Inside Out (DeeLIO 2022): The 3rd Workshop on Knowledge Extraction and Integration for Deep Learning Architectures_ , pages 100–114.
  * Lyu et al. (2022) Xinxi Lyu, Sewon Min, Iz Beltagy, Luke Zettlemoyer, and Hannaneh Hajishirzi. 2022\.  Z-icl: Zero-shot in-context learning with pseudo-demonstrations.  _arXiv preprint arXiv:2212.09865_.
  * Madaan et al. (2022) Aman Madaan, Niket Tandon, Peter Clark, and Yiming Yang. 2022.  Memory-assisted prompt editing to improve gpt-3 after deployment.  _arXiv preprint arXiv:2201.06009_.
  * Madaan and Yazdanbakhsh (2022) Aman Madaan and Amir Yazdanbakhsh. 2022.  Text and patterns: For effective chain of thought, it takes two to tango.  _arXiv preprint arXiv:2209.07686_.
  * Min et al. (2022) Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. 2022.  Rethinking the role of demonstrations: What makes in-context learning work?  _arXiv preprint arXiv:2202.12837_.
  * Mishra et al. (2022) Swaroop Mishra, Daniel Khashabi, Chitta Baral, and Hannaneh Hajishirzi. 2022.  Cross-task generalization via natural language crowdsourcing instructions.  In _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 3470–3487.
  * Nguyen et al. (2016) Tri Nguyen, Mir Rosenberg, Xia Song, Jianfeng Gao, Saurabh Tiwary, Rangan Majumder, and Li Deng. 2016.  Ms marco: A human generated machine reading comprehension dataset.  _choice_ , 2640:660.
  * Ni et al. (2021) Jianmo Ni, Chen Qu, Jing Lu, Zhuyun Dai, Gustavo Hernández Ábrego, Ji Ma, Vincent Y Zhao, Yi Luan, Keith B Hall, Ming-Wei Chang, et al. 2021.  Large dual encoders are generalizable retrievers.  _arXiv preprint arXiv:2112.07899_.
  * Nie et al. (2022) Ercong Nie, Sheng Liang, Helmut Schmid, and Hinrich Schütze. 2022.  Cross-lingual retrieval augmented prompt for low-resource languages.  _ArXiv_ , abs/2212.09651.
  * Ram et al. (2023) Ori Ram, Yoav Levine, Itay Dalmedigos, Dor Muhlgay, Amnon Shashua, Kevin Leyton-Brown, and Yoav Shoham. 2023.  In-context retrieval-augmented language models.  _arXiv preprint arXiv:2302.00083_.
  * Robertson et al. (2009) Stephen Robertson, Hugo Zaragoza, et al. 2009.  The probabilistic relevance framework: Bm25 and beyond.  _Foundations and Trends® in Information Retrieval_ , 3(4):333–389.
  * Rubin et al. (2022) Ohad Rubin, Jonathan Herzig, and Jonathan Berant. 2022.  Learning to retrieve prompts for in-context learning.  In _Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_ , pages 2655–2671.
  * Suzgun et al. (2022) Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoc V Le, Ed H Chi, Denny Zhou, et al. 2022\.  Challenging big-bench tasks and whether chain-of-thought can solve them.  _arXiv preprint arXiv:2210.09261_.
  * Wang et al. (2022) Boshi Wang, Sewon Min, Xiang Deng, Jiaming Shen, You Wu, Luke Zettlemoyer, and Huan Sun. 2022.  Towards understanding chain-of-thought prompting: An empirical study of what matters.  _arXiv preprint arXiv:2212.10001_.
  * Wei et al. (2022a) Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. 2022a.  Emergent abilities of large language models.  _Transactions on Machine Learning Research_.
  * Wei et al. (2022b) Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoc Le, and Denny Zhou. 2022b.  Chain of thought prompting elicits reasoning in large language models.  _arXiv preprint arXiv:2201.11903_.
  * Xie et al. (2021) Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. 2021.  An explanation of in-context learning as implicit bayesian inference.  In _International Conference on Learning Representations_.
  * Ye et al. (2023) Jiacheng Ye, Zhiyong Wu, Jiangtao Feng, Tao Yu, and Lingpeng Kong. 2023.  Compositional exemplars for in-context learning.  _arXiv preprint arXiv:2302.05698_.
  * Ying (2019) Xue Ying. 2019.  An overview of overfitting and its solutions.  In _Journal of physics: Conference series_ , volume 1168, page 022022\. IOP Publishing.
  * Zhao et al. (2021) Zihao Zhao, Eric Wallace, Shi Feng, Dan Klein, and Sameer Singh. 2021.  Calibrate before use: Improving few-shot performance of language models.  In _International Conference on Machine Learning_ , pages 12697–12706. PMLR.

##  Appendix A Examples of Retrieved Demonstrations

Question |  BM25 Demo |  GTR Demo
---|---|---
Q: when does the new episodes of supernatural start?
A: October 12, 2017 |  Q: when does the new episodes of ghost adventures start?
A: June 16, 2018 |  Q: when does the next episode of supernatural come out?
A: April 5, 2018
Kaj Birket-Smith (20 January 1893 – 28 October 1977) was a Danish philologist
and anthropologist. He specialized in studying the habits and language of the
Inuit and Eyak. He was a member of Knud Rasmussen’s 1921 Thule expedition. In
1940, he became director of the Ethnographic Department of the National Museum
of Denmark.
question: Kaj Birket-Smith would have been a ripe old age of 128 if he were
still alive today. Is it true, false, or neither?
answer: false |  Kaj Birket-Smith (20 January 1893 – 28 October 1977) was a Danish philologist and anthropologist. He specialized in studying the habits and language of the Inuit and Eyak. He was a member of Knud Rasmussen’s 1921 Thule expedition. In 1940, he became director of the Ethnographic Department of the National Museum of Denmark.
question: Kaj Birket-Smith was on the Thule expedition. Is it true, false, or
neither?
answer: true |  Kaj Birket-Smith (20 January 1893 – 28 October 1977) was a Danish philologist and anthropologist. He specialized in studying the habits and language of the Inuit and Eyak. He was a member of Knud Rasmussen’s 1921 Thule expedition. In 1940, he became director of the Ethnographic Department of the National Museum of Denmark.
question: Kaj Birket-Smith was a very educated man about many different
cultures and expressed love in his field of expertise. Is it true, false, or
neither?
answer: neither
Q: The original retail price of an appliance was 60 percent more than its
wholesale cost. If the appliance was actually sold for 20 percent less than
the original retail price, then it was sold for what percent more than its
wholesale cost?
Options: (A) 20(B) 28(C) 36(D) 40(E) 42Step-by-step reasoning process:
wholesale cost = 100; original price = 100*1.6 = 160; actual price = 160*0.8 =
128.
A: (B) |  Q: A retail appliance store priced a video recorder at 20 percent above the wholesale cost of $200. If a store employee applied the 20 percent employee discount to the retail price to buy the recorder, how much did the employee pay for the recorder?
Options: (A) $198 (B) $216 (C) $192 (D) $230 (E) $240
Step-by-step reasoning process: Wholesale cost of video recorder = 200 $ Video
recorder was priced at 20 percent above 200 = 240 $ % discount given by store
employee = 20 Emlpoyee paid = .8 * 240 = 192 $
A: (C) |  Q: A retailer bought a machine at a wholesale price of $108 and later on sold it after a 10% discount of the retail price. If the retailer made a profit equivalent to 20% of the whole price, what is the retail price of the machine?
Options: (A) 81 (B) 100 (C) 120 (D) 135 (E) 144
Step-by-step reasoning process: My solution: Wholesale Price= 108 Retail
Price, be = x He provides 10 % discount on Retail price= x-10 x/100 This
Retail price = 20 % profit on Wholesale price x-10 x/100 = 108+ 1/5(108)
x=144;
A: (E)
Q: Lori wants to buy a $320.00 pair of shoes and a matching belt that is
$32.00. Her part-time job pays her $8.00 an hour. How many hours will she have
to work before she can make her purchase?
Step-by-step reasoning process: b"She wants to buy a pair of shoes for $320.00
and a belt for $32.00 for a total of 320+32 = $<<320+32=352.00>>352.00 Her
purchase will total $352.00 and she makes $8.00 at her part-time job so she’ll
have to work 352/8 = <<352/8=44>>44 hours
A: 44 |  Q: Joanne makes $16.00 working at her main job for 8 hours a day. She has a part-time job, working an extra 2 hours a day where she makes $13.50 an hour. How much money does she make if she works this schedule 5 days a week?
Step-by-step reasoning process: She works 8 hours a day at $16.00 an hour so
she makes 8 * 16 = $128.00 a day. She works this job 5 days a week so she
makes 128 * 5 = $640.00 in 5 days. She works 2 hours a day at $13.50 an hour
so she makes 2 * 13.50 = $27.00 a day. She works this job 5 days a week so she
makes 27 * 5 = $135.00. She makes $640 at her main job and $135 at her part -
time job so all total she makes 640 + 135 = $775.00 in 5 days.
A: 775 |  Q: Janice has been working part-time at a convenience store 5 days a week. She can earn $30 per day and can earn $15 more when she works a 2 hour overtime shift. If she works three overtime shifts this week, how much will she earn this week?
Step-by-step reasoning process: Janice can earn $30 x 5 = $150 per week. She
will earn $15 x 3 = $45 more if she works three overtime shifts. Therefore,
Janice will earn $150 + $45 = $195 this week.
A: 195
Q: If it socially acceptable to wear an icon depicting crucifixion?
Step-by-step reasoning process: The crucifixion of Jesus is a common sign used
by Catholics and Christian denominations.Many jewelry stores offer necklaces
with the Crucifixion of Jesus Christ.
A: yes |  Q: Was the Donatello crucifix identified in 2020 life size?
Step-by-step reasoning process: The crucifix discovered in the church of
Sant’Angelo depicts an adult man. The crucifix discovered in the church of
Sant’Angelo is 89 cm high. The crucifix discovered in the church of
Sant’Angelo was identified as being a work of Donatello. The average height of
an adult man has been at least 150 cm in historical times.
A: no |  Q: Did any cultures associate celery with death?
Step-by-step reasoning process: Ancient Greeks used garlands of celery leafs
to bury their dead. Ancient Greece was considered a culture.
A: yes
Table 4:  Examples of retrieved demonstrations. NQ, ANLI(r3), AQUQ, GSM8K,
StrategyQA


# DGR: Tackling Drifted and Correlated Noise in Quantum Error Correction via
Decoding Graph Re-weighting

Hanrui Wang*1, Pengyu Liu*2, Yilian Liu3, Jiaqi Gu4, Jonathan Baker5, Frederic
T. Chong6, Song Han1 1Massachusetts Institute of Technology 2Carnegie Mellon
University 3Cornell University
4Arizona State University, 5University of Texas at Austin, 6University of
Chicago

###### Abstract

Quantum hardware suffers from high error rates and noise, which makes directly
running applications on them ineffective. Quantum Error Correction (QEC) is a
critical technique towards fault tolerance which encodes the quantum
information distributively in multiple data qubits and uses syndrome qubits to
check parity. A classical decoder is needed to repeatedly process the syndrome
qubit information and identify the errors. Minimum-Weight-Perfect-Matching
(MWPM) is a popular QEC decoder that takes the syndromes as input and finds
the matchings between syndromes that infer the errors. However, there are two
paramount challenges for MWPM decoders. First, as noise in real quantum
systems can drift over time, there is a potential misalignment with the
decoding graph’s initial weights, leading to a severe performance degradation
in the logical error rates. Second, while the MWPM decoder addresses
independent errors, it falls short when encountering correlated errors typical
on real hardware, such as those in the 2Q depolarizing channel. Overlooking
such correlations can adversely impact the performance of the decoder.

To tackle these two challenges, we propose DGR, an efficient decoding graph
edge re-weighting strategy with no quantum overhead. It leverages the insight
that the statistics of matchings across decoding iterations offer rich
information about errors on real quantum hardware. By counting the occurrences
of edges and edge pairs in decoded matchings, we can statistically estimate
the up-to-date probabilities of each edge and the correlations between them.
The reweighting process includes two vital steps: alignment re-weighting and
correlation re-weighting. The former updates the MWPM weights based on
statistics to align with actual noise, and the latter adjusts the weight
considering edge correlations. These adaptations pave the way for a more
accurate and resilient quantum error correction pipeline.

Extensive evaluations on surface code and honeycomb code under various
settings show that DGR reduces the logical error rate by 3.6$\times$ on
average-case noise mismatch with exceeding 5000$\times$ improvement under
worst-case mismatch.

##  I Introduction

Figure 1: Compare traditional decoding method with our proposed decoding graph
re-weighting (DGR). DGR leverages decoding statistics to dynamically re-weight
the decoding graph for high-fidelity quantum error correction. ††*Equal
Contributions.

Quantum Computing (QC) has garnered substantial research interest as an
emergent computational model designed to address problems previously deemed
unsolvable with enhanced efficiency. A multitude of sectors and academic
disciplines stand to gain from the potentialities of QC, notably cryptography
[56], database search [29], combinatorial optimization [21, 42], molecular
dynamics [50], and machine learning [44, 43, 61, 62, 63, 64, 67, 71, 10, 41]
applications, etc.

Despite the exciting advancements, the qubits and quantum gates on current
quantum machines suffer from high error rates of $10^{-3}$ to $10^{-2}$,
preventing us from executing applications that demand significantly lower
error rates (below $10^{-10}$) [40, 37, 25]. Therefore, reducing quantum error
is a pressing demand to close the gap. Quantum Error Correction (QEC), an
essential solution to this challenge, lowers the error rate by integrating
redundancy, a process where the information from a single logical qubit is
distributed across multiple physical qubits (data qubits). Syndrome qubits
then perform iterative checks of the data qubits and provide the parity of
them in classical bits, which are then decoded by a classical decoder to infer
the error that occurs on data qubits. By increasing the redundancy, the
logical error rate plummets exponentially.

An effective decoder is critical for QEC to lower the error rate. Minimum-
Weight-Perfect-Matching (MWPM) is a promising candidate with high accuracy and
almost-linear complexity [23, 34]. In MWPM, a decoding graph is constructed
based on the selected code and noise model. Each node in the graph represents
a syndrome qubit, and each edge represents an error that can cause a syndrome
flip event on its connected nodes. The weight on each edge ($w$) is computed
from the error rate of the corresponding error ($p$) as
$w=-\log{\frac{p}{1-p}}\approx-\log{p}$. To deal with the degeneracy between
syndrome and errors, i.e., the same syndrome pattern can be caused by
different errors, the MWPM finds the most likely error, which corresponds to
the overall minimum weights of the selected error edges. Therefore, a precise
noise model is crucial for the high accuracy of the decoder.

However, MWPM decoders suffer from the inaccurate modeling of the real
hardware noises that ignores the temporal drift of the noise and correlation
between noises.

➊ Noise Drift-Unaware – First, the noise on real quantum hardware can suffer
significant drifts overtime [53] due to various causes such as thermal
fluctuations in Josephson Junctions [30], laser detuning in neutral atom and
trapped-ion [14], unwanted coupling to two-level-systems (TLS) [47] defects in
Transmon qubits [7, 55], magnetic fluctuation caused by the sun [46], etc. The
drifted noise creates a mismatch between the actual noise characteristics on
quantum hardware and the noise model inside the classical MWPM decoder. Though
there exist traditional methods for identifying noise in quantum systems,
e.g., quantum process tomography [11, 51], randomized benchmarking [20, 39,
45], they are time-consuming and cannot be performed simultaneously with QEC
in progress.

➋ Noise Correlation-Unaware – Another issue is the lack of correlation
modeling between noises (represented as edges in the decoding graph).
Typically, edges within this graph are treated as independent. Yet, in many
cases, these edges exhibit nontrivial correlations. As an illustrative
example, consider the widely-referenced X-Z surface code. Here, a Y error in
the 1Q depolarizing channel decomposes into two ostensibly independent edges:
one X and one Z. However, when a Y error occurs in practice, both edges should
be co-selected. Therefore, the upper-bound accuracy of MWPM is unfortunately
compromised due to its neglect of edge correlations.

To overcome the above two issues and augment the existing MWPM pipeline with
more accurate noise modeling, we propose an efficient, non-intrusive, and
scalable approach – Decoding-Graph Re-weighting (DGR) (Fig. 1). DGR operates
based on two pivotal insights. The first recognizes that quantum error
correction inherently possesses the ability to detect errors, thus allowing
the reconstruction of noise characteristics from multiple rounds of QEC. The
second insight acknowledges that the method can discern not only independent
noise but also the correlations between noises. This correlation information
can be harnessed to fine-tune the weights on the graph edges, thereby offering
a more robust solution.

Figure 2: Surface code and its decoding graph.

To address the noise drift issue, we propose an alignment re-weighting
strategy. In QEC, the syndrome extraction process is performed iteratively
until the whole quantum program is completed. That provides massive
information about the errors on quantum hardware. Take the common standard for
real-time decoding [2], 1 $\mu$s per round, 1 million rounds of matchings can
be obtained every second. The occurrence of an edge in the matching indicated
a predicted error on that edge. Therefore, we first leverage an occurrence
tracer to trace the frequency of each edge that occurred in the matchings.
Given enough rounds, we can accurately estimate the correct weight for every
edge. Then, the MWPM can leverage the updated weights in the later rounds to
improve accuracy. This re-weighting is non-intrusive, running in parallel on
the classical sides, with low overhead. For instance, with a distance of 5 and
a physical error rate of 0.05, 10,000 rounds are enough to have an accurate
estimation.

To further consider noise correlation, we propose correlation re-weighting. A
correlation tracer is used to count the frequency of co-occurrence of every
edge pair and construct a correlation matrix. Then, we adopt a two-iteration
MWPM in this strategy. In the first iteration, we perform a normal MWPM, and
then, according to the edges in the initial matching, we slightly adjust the
weights of each edge based on the correlation matrix. We will increase the
prob of an edge if its correlated edges are selected in the initial mapping.
To give a simple example, if a 1Q X edge is selected, then we will increase
the correlated Z edge because the X error may be due to a Y error. Then, we
will use the re-weighted graph to perform the second iteration MWPM and obtain
the final matchings. We provide two re-weighter options, one heuristic-based
and one neural network (NN)-based, to determine the new weights of the edges
according to correlations. To reduce the overhead, we only perform this re-
weighting for difficult syndrome patterns, such as the ones with a relatively
large number of syndromes.

DGR is a plug-and-play scheme that can be combined with a variety of existing
optimizations on MWPM. More importantly, it incurs zero quantum overhead since
there is no additional operation on quantum hardware.

We extensively evaluate DGR on over 40 benchmarks with Surface Code and
Honeycomb Code under different code distances, error rates, and noise models
(phenomenological and circuit-level). DGR brings more logical error rate
reduction with a larger code distance and a lower physical error rate, with
3.6$\times$ and 1.7$\times$ average-case reduction for surface code and
honeycomb code. For the worst-case mismatch, the reduction for surface code is
695$\times$ on average, up to 7,360$\times$. The correlation re-weighting
further reduces the error rate by an average of 1.2$\times$ and up to
1.4$\times$. We also perform sensitivity analysis on the effectiveness of DGR
and observe its superior capability to recover correct weights even under a
500$\times$ error rate drift. The key insight we derive is that, even if the
mismatch causes a much higher logical error rate, its absolute value remains
much smaller than 1, such that most of the matchings, especially for simple
cases, are still correct and can be used to recover the correct weights.

Our contributions are summarized as follows:

  * •

We show that the mismatch between the noise model in MWPM and real hardware
noise results in a drastic degradation of logical error rates by several
orders of magnitude.

  * •

To augment the current MWPM flow with more accurate noise modeling, we propose
an efficient, non-intrusive, and scalable approach, DGR, with zero quantum
overhead.

  * •

We propose alignment re-weighting to trace the frequency of edges from
matchings and update edge weights. The process can fully recover the logical
error rates to reach similar rates to a weight oracle.

  * •

We propose a correlation re-weighting scheme that first traces the
correlations between edge pairs based on matchings and then adjusts the edge
weights with a simple heuristic or an NN-based re-weighter.

  * •

We evaluate DGR on surface code and honeycomb code under various settings and
achieve 3.6$\times$ and 1.7$\times$ average-case improvement. The worst-case
improvement can exceed 5000$\times$ for surface code.

##  II Background and Related Work

###  II-A Surface Code and Honeycomb Code

The Surface code is a prominent QEC scheme that encodes a logical qubit into a
two-dimensional lattice of alternating data and syndrome (parity) qubits. It
has a high error threshold and requires only nearest-neighbor connectivity,
making it practical for real quantum systems. Errors on data qubits are
categorized into a discrete set of Pauli errors - bit-flip (Pauli-X), phase-
flip (Pauli-Z), or both [49]. They can be detected by adjacent parity qubits
using a stabilizer circuit, which measures a four-qubit operator, leading to
the detection of X, Z, or Y (combination of X and Z) errors. The surface code
can correct error chains up to length $\lfloor\frac{d-1}{2}\rfloor$. A simpler
variant of the original Toric code [36], the ‘rotated’ surface code as in 2
left is often preferred due to its more compact layout, reducing the physical
qubits and gate overheads. The $[[d^{2},1,d]]$ stabilizer code has become a
prime candidate for near-term fault-tolerant quantum computation.

Though the surface code has many merits, it requires connectivity of 4. It
could be directly mapped to some superconducting hardware [1] and neutral atom
devices [68, 66] which means each qubit needs to connect with 4 other qubits,
but may be difficult for some others. For example, IBM’s newest quantum
computer employs a heavy-hexagon lattice, which provides connectivity of only
2.4 [9]. Such sparse connectivity challenges the design of quantum error
correction code. Recently, Honeycomb code [33, 32], a dynamically generated
quantum error correction code, is proposed to perform error correction on
quantum computers with heavy-hexagon topology. As a instantaneous stabilizer
code, it shares many similar features to conventional stabilizer code: in each
step, the code can be viewed as stabilizer code, but the stabilizer group
changes over time. An advantage of honeycomb code is that its decoding can
also be handled with MWPM [27, 26]. It is worth noticing that the measurement
error, single qubit error, and two qubits error will all be correlated edges
in the decoding graph. Thus, decoding with an MWPM decoder is not optimal for
honeycomb.

###  II-B Minimum-Weight-Perfect-Matching

To decode errors from syndromes, researchers have proposed multiple families
of decoders such as Lookup Table (LUT) based [12, 57, 54], Neural Network (NN)
based [58, 6, 8, 65] and Union-Find (UF) based [16, 15, 69, 13]. Among various
decoders, Minimum-Weight-Perfect-Matching decoders [19, 18, 22, 35] are
considered to have a good balance between the decoding accuracy and speed. It
has almost linear time complexity [23, 34, 70], is practical for hardware
implementation for real-time decoding [59] and is more accurate than LUT and
Union Find decoders.

MWPM constructs a decoding graph from the code as in Fig. 2. Each node
corresponds to a syndrome qubit. Each edge represents an error that can cause
the syndrome to flip on its connected nodes. Most edges have two nodes, and
some connect one node to the boundary. In reality, the syndrome extraction
circuit and the syndrome measurements also have errors, so the decoder
combines multiple rounds of the syndrome (typically $d$). This is referred to
as the circuit-level noise model, where the decoding graph is a more
complicated 3D graph with edges connecting nodes between adjacent rounds.

Each edge also has a weight $w=-\log{\frac{p}{1-p}}$. The error probability is
obtained from a noise model. In MWPM assumption, all errors happen
independently, and each error will only cause one or two detection events
corresponding to one edge in the matching graph. Thus, the highest probability
error event can be found as the minimum weight matching of the decoding graph.
However, on real quantum devices, not all errors happen independently, and in
many realistic QEC codes, each error may cause more than two detection events.
For example, for surface code, a Y error on the data qubit will flip the four
syndrome qubits around it, creating four detection events. Besides, since the
same syndrome can be caused by different error events, a precise alignment
between the decoder’s input noise model and real hardware’s noise
characteristics is crucial for accurate QEC decoding.

##  III Motivation

Figure 3: Noise drift with in short-term and long-term.

###  III-A Noise Drift

The effectiveness of the decoder is inherently linked to the underlying noise
model of the quantum device. Noise models characterize the physical processes
that can cause errors, such as bit flips and phase flips, and they provide
essential information to compute the probability of particular error patterns.

Unfortunately, noise in quantum systems is known to undergo significant drift
over time [53], which can arise from various sources. For superconducting
quantum devices, thermal fluctuations can be caused by instantaneous heat
released by a quantum phase slip [30] in Josephson Junction, leading to noise
drift. Unwanted coupling between the two-level-systems (TLS) [47] defects and
Transmon qubits [7, 55] can cause the fluctuation of T1 and T2 coherence
times. The population of unpaired electrons can also lead to decoherence time
changes [31]. In neutral atom and trapped-ion systems, photons can be excited
from electric-field noise, varying the gate fidelity [3]. Laser instabilities
such as detuning [14] can also drift the fidelity of gates. Besides, the
magnetic field of the Earth is subject to fluctuations influenced by solar
phenomena, including the incremental changes that follow the Sun’s daily
patterns and the more abrupt alterations stemming from events like coronal
mass ejections. The variations in the magnetic field can cause fluctuations in
qubit frequencies [46]. Cross-talk between qubits, quasi-particles [38], and
seismic noise from human activities are also responsible for noise drift.

Fig. 3 shows our observations on noise fluctuation in IBM and Regetti
superconducting devices. Subfig (a) shows the drift of 1Q and 2Q gate error
rates within a short period – 100 seconds. The y-axis is the error rate ratio
normalized by the smallest error rate. The drift can exceed 15$\times$ within
only 100 seconds. Subfig (b) shows the historical error rate data extracted
from IBM devices. We show the distribution of the max drift of the error rate
of a gate in the past year. The 1Q and 2Q gate drifts can easily exceed
10$\times$. 1Q gate drift is larger than 2Q gate, and can be over
1000$\times$.

To accurately measure and characterize the noise profile, techniques such as
quantum process tomography [11, 51], gate set tomography [48, 28], randomized
benchmarking [20, 39, 45], and cross-entropy benchmarking [5] can be employed.
The accuracy comes at the cost of inefficiency. These techniques are simply
not scalable as the quantum systems grow in complexity. The scalability issue
prevents them from being used for noise models that update in real-time.
Besides, the techniques above are also impractical to be performed together
with the QEC. Although some prior work proposed to capture the Pauli error
[60] from syndromes, we pursue to cover arbitrary type of errors (Pauli,
measurement, etc.) and capture the correlations between different errors.

Figure 4: Re-decoding with correlation re-weighting to identify the correct
errors. Figure 5: The overview of the proposed decoding graph re-weighting
flow, including conditional correlation re-weighting to capture edge
correlation and periodic alignment re-weighting to align the probability.
Figure 6: Proposed alignment re-weighting method. Based on multiple matching
trials, the occurrence tracer and correlation tracer will count the single-
edge and edge-pair probability, respectively. Based on the occurrence tracing
result, the alignment re-weighter will assign the negative log probability as
the edge weight. The correlation tracer is used in the correlation re-
weighting module.

###  III-B Noise Correlation

The MWPM decoder has the assumption of independence between all the edges in
the decoding graph. However, in reality, the edges are correlated. The edge
correlations can come from two sources. In the one-to-many type of
correlation, one single error can impact multiple edges due to the code
construction. For instance, the Y error causes errors on one X and Z edge in
surface code; an error on the X-basis measurement of a syndrome qubit causes
errors on one Y edge and one Z edge in honeycomb. The second is the many-to-
many type of correlation, in which multiple errors happen together and cause
correlation in their edges. For example, the 2Q depolarizing channel can cause
two Pauli errors on two qubits. Besides, for neutral atom QEC [4], one laser
controls a whole row or column of qubits, whose error can also be correlated.

The correlation information can be leveraged to augment the MWPM decoder. For
example, in Fig. 4, the actual errors in (a) contain one X and one Y, which
cause three syndromes in (b). With the initial decoding round, the X error
matching (blue) will match the syndrome with the bottom boundary because of
the lower weight. There are two syndromes for Z error matching (yellow), and
the results indicate a Z error. If we know a Y error is causing that Z, we can
lower the weight of the corresponding X error edge. The re-decoding process
runs MWPM again and can match the syndrome with the top boundary, which is the
correct decoding. In reality, the edge correlations can be obtained from
decoding history. Some rule-based or data-driven strategies can be developed
to perform the re-weighting based on initial matching and correlations.

##  IV Methodology

An overview of DGR workflow is shown in Fig. 5. It augments the conventional
decoding pipeline by performing alignment re-weighting and correlation re-
weighting. The process of alignment re-weighting within the quantum error
correction system may be executed either periodically or as necessitated by
the drift of noise. It will look back to extract a specific number of
immediate preceding decoding trials and leverage the decoded matching from
those trials to compute updated weights. These recalculated weights, being in
alignment with the existing real noise on the quantum hardware, are then
implemented in subsequent MWPM decoding processes. Moreover, an additional re-
weighting step, correlation re-weighting, is strategically employed on what is
characterized as “difficult cases” to find a more probable error path when
taking into consideration the correlation of edges within the system.

###  IV-A Alignment Re-weighting

The alignment re-weighting is shown in Fig. 6 (a). It contains an occurrence
tracer and a re-weighter. The occurrence tracer keeps track of the previous
decoding results. When alignment re-weighting is needed, the occurrence tracer
estimates the occurrence frequency of each edge in previous rounds and uses
this frequency as the probability to update the weight in the MWPM decoding
graph. In the figure, we show specific examples. The red circles mark the
syndrome locations. Thick edges are the matched edges generated by the MWPM
decoder. The tracer looks back at a set number of decoding trials and keeps a
counter for each edge. The counter goes up by one every time that edge appears
in the matching.

It is clear that the accuracy of the estimated weight is bounded by the
decoder’s performance. A poor decoder cannot provide reliable decoder
matchings. However, our sensitivity analysis in Section V proves that the re-
weighting still works even if the actual error rate drifts by 100$\times$. Our
key insight here is that most syndromes are isolated when the physical error
rate is below a certain threshold. For those simple isolated syndromes, the
decoding accuracy is robust to the changes in edge weights. Let’s consider two
syndromes caused by a single, isolated error, as shown in Fig. 6 (b). The most
straightforward matching is directly matching those two syndromes with one
single edge $e_{3}$. The second shortest path between the syndromes has three
edges (another three edges of a square $e_{8},e_{9},e_{10}$). Since the weight
$w$ is the negative logarithm of the probability $p$ (expressed as $w\approx-
log(p)$), the error rate $p$ would have to drift roughly three orders of
magnitude to cause a mistake in matching. Therefore, our method can handle a
wide range of noise drift. If the error is not isolated, the chance of one
error in a fixed area is $O(p_{phys})$, and more than two errors is
$O(p_{phys}^{2})$. As long as $p_{phys}^{2}\ll p_{phys}$, we can estimate the
error with high accuracy and robustness.

Figure 7: Weight estimation error for surface code with distance $5$ and
honeycomb code with distance $3$ under circuit level noise $p=0.001$. Figure
8: Heat map of the correlation strength for surface code with distance 3 under
phenomenlogical and circuit level noise and honeycomb code with distance 3
under circuit level noise. The lighter the point is, the higher the
correlation is.

The number of trials required to statistically estimate the weight is
influenced by the physical error rate; it increases as the error rate
decreases but remains independent of the code distance. As illustrated in Fig.
7, the mean squared error of estimation tends to converge after $10^{6}$
trials for both surface code with distance $d=5$ and error rate $p=0.001$, and
honeycomb code with distance $d=3$ and error rate $p=0.001$.

Figure 9: Number of related edges for two codes.

###  IV-B Correlation Re-weighting

The second part of DGR is called correlation re-weighting. It is comprised of
two components: a correlation tracer and a correlation re-weighter. The
tracer’s function is to construct a correlation matrix by counting the
occurrence of edge pairs in previous trials. Though actual correlations may
involve multiple edges, here, we restrict our count to correlations solely
between two edges for two principal reasons. First, the memory consumption of
tracking correlations among multiple edges escalates exponentially. Second,
estimating multi-edge correlation is challenging and sometimes impossible.
Take the surface code as an example. Suppose the X-check
$X_{1}X_{2}X_{3}X_{4}$ forms a stabilizer. Then, the error $X_{1}$ and error
$X_{2}X_{3}X_{4}$ will generate exactly the same syndromes and can never be
distinguished through the measurement data. The ground truth correlation
matrices and the estimated ones for surface code and honeycomb code are shown
in Fig. 8. Most entries in the matrices are accurately estimated, with some
discrepancies stemming from the complexities of multi-edge correlations.

The MWPM decoder operates by performing certain approximations and views all
the edges independently, an approach that can impair its performance. To
enhance the decoding process by leveraging correlation information, we
introduce a two-iteration decoding methodology. In the first iteration,
standard MWPM decoding is executed. In the second, we adjust the weights based
on both the correlation matrix and the initial matching and re-decode the
syndromes. The key insight here is that the known occurrence of one edge
increases the probability of occurrence for its correlated edges. To limit the
overhead for two-iteration decoding, this technique is selectively applied to
challenging syndrome patterns (characterized by the number of syndromes),
which is inspired by the two-level decoding approach in Clique [52]. The
comprehensive illustration of the entire pipeline is shown in Fig. 10.

Figure 10: Proposed correlation re-weighting method. Once the correlation re-
weighting condition is triggered, the edge co-occurrence probability will be
obtained from the correlation tracer and fed to the correlation re-weighter,
either a heuristic or a NN-based re-weighter, for edge weight adjustment. And
then an additional MWPM is performed with adjusted weights.  Figure 11: The
proposed neural network-based correlation re-weighter trained with SPSA
algorithm.

####  IV-B1 Heuristic-Based Correlation Re-weighter

The core intuition behind correlation re-weighting lies in the relationship
between physical errors and their corresponding edges within the matching
graph. Specifically, if a single physical error leads to two or more edges in
this graph, a positive correlation emerges among these edges. Consequently, if
one edge is chosen in the predicted matching, the probability of the other
edges occurring should be increased, resulting in a decrease in their weights.
Conversely, if one edge is not selected, the probability of the occurrence of
the other correlated edges should be diminished, leading to an increase in
weights.

With regard to the correlation re-weighter, one straightforward heuristic is
to re-weight according to the conditional probability of edge matching in the
initial matching [17]. Unfortunately, this approach is viable only when an
edge has one single correlated edge. However, as in Fig. 9, after evaluating
the number of related edges on surface code with distance 5 and honeycomb code
with distance 3, it’s apparent that the majority of edges have more than one
correlated edge, so the application of a simple conditional probability
becomes impossible. This complexity necessitates a more nuanced approach to
accurately capture the correlations and their effects on the decoding process.
As a result, we propose a heuristic re-weighting algorithm that is able to
process multiple correlated edges, increasing or decreasing the edges’ weight
according to the strength of correlation.

After the first decoding round, we get the initially selected edge set $M$;
then, we update all the edges’ weights according to the following rule.

$\begin{split}\widetilde{w_{j}}&=w_{j}-\sum_{e_{i}\in
M}\frac{p(e_{i},e_{j})}{p(e_{i})}+\sum_{e_{i}\not\in
M}\frac{p(e_{i},e_{j})}{p(e_{i})}\\\ \end{split}$ (1)

where $w_{j}$ and $p(e_{i})$ represent the weight and occurrence probability
of edge $e_{j}$, $e_{i}$, respectively, as estimated by the occurrence tracer,
and $p(e_{i},e_{j})$ denotes the probability of edge $e_{i}$ and $e_{j}$ co-
occuring. Subsequently, we proceed with a second decoding round, utilizing
these newly calculated weights to refine the results.

Algorithm 1 Hybrid training for NN reweighing predictor

0: Decoding graph G, NN predictor $\Psi_{\Theta}(\cdot)$, Dataset
$\mathcal{D}$, Steps $T$, Learning rate $\eta$, initial weight $w$

1: for each training iteration $t\leftarrow 1\cdots T$ do

2: Sample a mini-batch $(x,m^{\prime})$ from $\mathcal{D}$

3: Forward the MLP predictor

4: $\widetilde{w}=w+\Psi_{\Theta^{t}}(x)$

5: Forward the MWPM decoder and calculate the loss

6:
$\mathcal{L}(\widetilde{w})=\mathcal{L}(\texttt{MWPM}(G(\widetilde{w})),m^{\prime})$

7: Estimate the gradient through the combinatorial MWPM decoder

8:
$\nabla_{\widetilde{w}}\mathcal{L}=\frac{1}{2Q\sigma^{2}}\sum_{i=1}^{Q}\big{(}\mathcal{L}(\widetilde{w}+\Delta
w_{i})-\mathcal{L}(\widetilde{w}-\Delta w_{i})\big{)}\Delta w_{i}$

9: Use the estimated gradient for backpropagation

10:
$\nabla_{\Theta^{t}}\mathcal{L}=\nabla_{\widetilde{w}}\mathcal{L}\cdot\nabla_{\Theta^{t}}\widetilde{w}$

11: Update the NN predictor weights

12: $\Theta^{t+1}\leftarrow\Theta^{t}-\eta\nabla_{\Theta}\mathcal{L}$

13: end for

14: return Trained NN predictor $\Psi_{\Theta^{*}}(\cdot)$

####  IV-B2 NN-based Correlation Re-weighter

Figure 12: Logical error rate with and without alignment re-weighting on
10$\times$ mismatch for surface code under phenomenological noise.

If the edge is associated with a multitude of related edges with complex
relationships, the heuristic-based re-weighting may not provide effective
results. To fully leverage the correlation matrix derived from the matching
history, we introduce a neural network-based re-weighter. This re-weighter is
designed to predict weights changes, denoted as
$\Delta\widetilde{w}(e_{i})=\Psi_{\Theta}(x(e_{i}))$, where $x(e_{i})$
symbolizes the extracted feature for edge $e_{i}$, as shown in Fig. 10 and
Fig. 11, and $\Psi_{\Theta}(\cdot)$ is the NN-based predictor. The weight for
the second iteration of decoding is calculated as
$\widetilde{w}(e_{i})=w(e_{i})+\Delta\widetilde{w}(e_{i})$. The extracted
feature contains information on the edge type, which is encoded as a one-hot
vector. The edge type is determined by the spatial location of an edge in the
decoder graph. The edges of the same type will have the same number of
correlated edges. Furthermore, the feature vector incorporates information on
correlated edges, including whether the related edge has been matched in the
initial iteration matching $M_{0}$ and the co-occurrence probability
$p(e_{i},e_{j})$. We make sure that the order of the correlated edges in the
vector consistently follows the same related location order for one edge type,
such as from front to back, from top to bottom, etc. This uniform ordering is
necessary to maintain the spatial invariance of the code. To accommodate
varying numbers of related edges, the feature vector is zero-padded to the
maximum number of related edges, ensuring a consistent input for training and
inference.

The complexity of model training represents the most demanding part of this
process, and it can be mathematically formulated as

$\min_{\Theta}~{}~{}\sum_{e_{i}\sim
E}\mathcal{L}\Big{(}\texttt{MWPM}\big{(}G(w+\Psi_{\Theta}(x(e_{i})))\big{)},m^{\prime}\Big{)},$
(2)

where the objective is to minimize the error between the ground-truth matching
solution $m^{\prime}$ and the solution on the re-weighted decoding graph with
the predicted edge changes of the weights $\Delta w$.

The gradients w.r.t. the NN weights are expressed
as$\nabla_{\Theta}\mathcal{L}=\nabla_{\widetilde{w}}\mathcal{L}\cdot\nabla_{\Theta}\widetilde{w}$.
A significant challenge arises from the combinatorial nature of the solver,
MWPM($\cdot$), which is not differentiable, i.e., the first term
$\nabla_{\widetilde{w}}\mathcal{L}$ cannot be obtained directly. To navigate
around this limitation and incorporate the black-box combinatorial solver into
the backpropagation procedure, we adopt a symmetric zeroth-order gradient
estimator to obtain the directional derivative, as described in Alg. 1 Lines
7-10. To achieve this, we randomly generate a perturbation of the weights from
a multi-variate Gaussian distribution $\Delta w\sim\mathcal{N}(0,\sigma^{2})$,
and use the directional derivative as an estimate of the first-order gradient
of the smoothed MWPM oracle. To reduce the sampling variance, we average the
gradients over $Q$ samples. Subsequently, the estimated gradients of the
solver are used to calculate the gradients of the weights of the NN model.

During the inference phase, the model can process all edges in a batch,
thereby predicting the $\Delta\widetilde{w}(e_{i})$ to minimize latency.

##  V Evaluation

###  V-A Evaluation Methodology

Benchmarks. We evaluate two quantum error correction codes, the surface and
honeycomb codes. For surface code with code distance $d$, we use stim’s [24]
built-in functions surface_code:rotated_memory_z to generate a surface code
with $d$ rounds and test with logical Z observable. For honeycomb code, we use
the package provided by Ref. [27] to generate a planar honeycomb code. For a
code distance with $d$, we use $d+1$ width, $3\times[(2d+2)/3]$ height and
$3d$ rounds, and test with a V-type observable. We use PyMatching [34] to
sample the errors and decode them. We select distances 3, 5, 7, and 9 for the
surface code with physical error rates from 2E-4 to 3E-2. We select distances
3 and 5 for the honeycomb code with a physical error rate from 3E-5 to 1E-3.

Figure 13: Logical error rate with and without alignment re-weighting on
10$\times$ mismatch for surface code under circuit-level noise.

Noise Model. We implement two types of noise model, the phenomenological noise
model [18], which contains data qubit depolarization and measurement error,
and the circuit-level noise model, which includes data qubit depolarization,
depolarization after Clifford gates, and measurement error. For surface code,
the errors are generated using stim’s built-in functions, and for honeycomb
code, we used the SD6 circuit level noise model introduced in [27].

Mismatch. We evaluate two mismatch models. The first one is random (average-
case mismatch). We first generate an error correction circuit with a given
physical error rate in the random mismatch model. Then, we mutate each
operation’s error probability with a factor, sampled log-uniformly from
$\frac{1}{N}$ to $N$. This is called $N$$\times$ mismatch. The circuit and the
matching are sampled independently. The second one is the worst-case mismatch,
in which the difference between the logical error rate using the matched
decoder and the mismatched one is maximized. We found that for surface code
with logical Z observable, the logical error rate will be maximized, if the
first $\frac{d-1}{2}$ rows of the data qubits have a larger physical error
rate and the last $\frac{d+1}{2}$ rows have a lower physical error rate.

NN Decoder Settings The NN we adopt is an MLP with two hidden layers with
dimension 64. The last layer is a tanh output layer scaled by 3$\times$,
outputting $[-3,3]$. We leverage the Adam optimizer with a learning rate of
1E-3 and weight decay of 1E-4. The models are trained with batch size 128 for
100 epochs on datasets with 100000 generated samples.

Figure 14: Logical error rate with and without alignment re-weighting on
10$\times$ mismatch for honeycomb code under phenomenological and circuit-
level noise. Figure 15: Logical error rate with and without matching
prediction on 10$\times$ mismatch. TABLE I: Physical error rate threshold for
surface code.

Mismatch Strength | Oracle | Mismatched | Ours
---|---|---|---
10$\times$ | 0.0301 | 0.0194 | 0.0300
5$\times$ | 0.0272 | 0.0234 | 0.0271

Figure 16: Logical error rate reduction with the proposed correlation re-
weighting scheme.

###  V-B Main Results

Logical error rates reduction with alignment re-weighting. Fig. 12 shows the
performance of surface code under phenomenological noise. The y-axis is the
relative logical error rate of decoder with mismatch weights and our aligned
weights, over the decoder with matched weight provided by an oracle. Our
aligned weights is obtained from matching of 1 million trials. The mismatch
strength is 10$\times$. Subfig. (a) shows the average performance over 10
random sampled noise model. The mismatched weights degrade the decoder by
3.6$\times$ on average, while our alignment re-weighter can fully recover the
logical error rate. Generally, we observe the improvements of our re-weighter
is more significant for large distances and low physical error rates. Subfig.
(b) shows the performance in the worst case mismatch where the mismatched
weights can cause logical error rate to degrade by up to 1582$\times$, with
77.5$\times$ on average. Fig. 13 further compares the mismatched weight with
our aligned weight under the circuit-level noise model. For random benchmarks,
DGR provides 1.7$\times$ improvements and under worst case mismatch, DGR
provides 695$\times$ reduction on average (764.3/1.1), up to 7360$\times$. Our
aligned weights cannot fully recover performance, leading to a 1.1$\times$
higher logical error rate. Fig. 14 provides more results on the honeycomb
code. Again, DGR reduces logical error rate by 1.6$\times$ and 1.8$\times$ for
two noise models and fully recovers them to the level of using weight oracles.
From these results, we demonstrate the detrimental effect of mismatched
weights, and prove the effectiveness of our DGR on various code types,
distances and error rates.

Figure 17: Performance of alignment re-weighter v.s. different mismatch level.
Evaluated on surface code with distance 5 under phenomenological noise with
$p=0.001$. Figure 18: Logical error rate vs. the number of trials under
different noise level, evaluated on surface code under phenomenological noise.
Figure 19: Logical error rate vs. the number of trials with different code
distances, evaluated on surface code under phenomenological noise. TABLE II:
Required memory size for storing 10000 rounds of matching trials on classical
machine.

Memory (KB) | p=1E-3 | =p2E-3 | p=5E-3 | p=1E-2 | p=2E-2 | p=3E-2
---|---|---|---|---|---|---
d=3 | 0.39 | 0.78 | 1.95 | 3.9 | 7.8 | 11.7
d=5 | 2.51 | 5.02 | 12.55 | 25.1 | 50.2 | 75.3
d=7 | 7.79 | 15.58 | 38.95 | 77.9 | 155.8 | 233.7
d=9 | 17.67 | 35.34 | 88.35 | 176.7 | 353.4 | 530.1

TABLE III: Number of edges, feature dimension and latency of NN re-weighter
for different QEC.

Code | surface | surface | surface | surface | honeycomb | honeycomb
---|---|---|---|---|---|---
Distance | 3 | 5 | 3 | 5 | 3 | 5
Error Model | pheno | pheno | circ | circ | circ | circ
N_edges | 59 | 313 | 86 | 518 | 441 | 2499
N_features | 15 | 18 | 80 | 139 | 148 | 221
Latency (ms) | 0.14 | 0.23 | 0.15 | 0.31 | 0.28 | 0.77

Logical error curve. Fig. 15 shows the logical error rate on surface code with
distance 5 on the phenomenological noise model. As the physical error rate
increases, the ratio between the mismatched logical error rate and the matched
one decreases. This is because both of them will converge to 0.5 when the
physical error rate is large enough. The trend on the worst case mismatch is
more significant. When the physical error rate increases, our method’s
performance slightly decreases because more errors happen in the decoding
procedure, thus making the estimation of weights inaccurate.

Physical error threshold. Tab. I shows the physical error threshold for
surface code. In both 10$\times$ and 5$\times$ mismatch strength, the
threshold of the mismatched decoder is significantly lower, and the predicted
weights have a performance comparable to the ideal weight oracle.

Logical error rates reduction with correlation re-weighting. Fig. 16 (a) and
(b) show the logical error rate reduction for surface code with heuristic and
NN correlation re-weighter on two different noise models. The heuristic re-
weighter reduces the logical error by 1.01$\times$ and 1.05$\times$ on average
for surface code and 10$\times$ Y-biased surface code, while the NN re-
weighter manages to reduce by 1.11$\times$ and 1.17$\times$. Fig. 16 (c) shows
the logical error rate reduction for honeycomb code with heuristic and NN re-
weighter on the phenomenological and circuit level noise models. The heuristic
re-weighter reduces the logical error by 1.15$\times$ on average, while the NN
re-weighter reduces the logical error by 1.12$\times$. Although the heuristic-
based re-weighter sometimes works better than the NN-based re-weighter. The
NN-based one has a more stable performance and is capable of handling
different types of QEC codes under diverse settings.

###  V-C Analysis

How severe mismatch can our method recover? Fig. 17 shows the performance of
alignment re-weighter when increasing the mismatch level. For the random case,
the alignment re-weighter can always recover the correct weight, up to
500$\times$ mismatch. In the worst case, the alignment re-weighter can recover
the weight almost perfectly when the mismatch is smaller than 50$\times$.
After that, the logical error rate degrades when the mismatch is more severe,
but it is still better than the mismatched decoder. As analyzed in Section
IV-A, we expect the weight estimation will be less accurate when a direct edge
between two syndrome nodes could be mis-decoded to a three-edge detour. That
will happen when the mismatch strength is 1000$\times$, which is close to our
observations here for the average-case mismatch.

Required trials for accurate estimation vs. physical error rate. Fig. 18 shows
the convergence of logical error rates under different noise levels. The
smaller the noise is, the more trials are required. This is because whether an
edge is selected forms a Poisson distribution with expectation $p$. Estimating
the expectation value up to a constant relative error requires trials
inversely proportional to $p$.

Required trials for accurate estimation vs. code distance. Fig. 19 shows the
convergence of logical error rate with the increasing of trials. For different
code distances, the required trial is similar, around $10^{3}$ to $10^{4}$.
This is because, under the same error rate, the error occurrence of an edge is
not influenced by the size of the decoding graph.

Memory Requirements. Table II shows the additional memory requirement for the
occurrence and correlation tracers to work. We use a sparse matrix with INT16
data type to store previous decoding results up to N trials. Thus, whenever
needed, the two tracers can estimate the weights and correlations from the
data. For the most resource-consuming case, $d=9$ and $p_{phys}=0.03$, the
memory requirement is only 0.5MB. The memory requirement can be further
decreased by storing more coarse-grained data.

Overhead of NN-reweighter. Table III shows the overhead of the NN re-weighter.
The latency ranges from 0.14ms to 0.77ms according to the scale of the
problem. We can further reduce the average latency by only choosing the
difficult cases with more syndromes to go through the NN re-weighter. In our
experiments above, only 15% of trials need correlation re-weighting. The easy
cases with fewer syndromes can be correctly decoded in the first iteration
with high probability.

##  VI Conclusion and Outlook

We propose DGR, an efficient decoding graph edge re-weighting strategy to
solve the noise drift-unawareness and correlation-unawareness of the MWPM
decoder. DGR encapsulates alignment re-weighting that statistically estimates
the probabilities of individual edge according to the history of decoded
matchings and update the edge weights. Besides, correlation re-weighting takes
the correlation between edges into consideration to adjust the edge weight and
perform re-decoding. We propose one heuristic-based re-weighter and one neural
network-based re-weighter, and both of them can improve the logical error rate
leveraging correlations. Looking forward, we expect the noise drifts and
correlations on large-scale quantum systems will likely be more complicated.
We hope DGR provides a perspective on dealing with them with small classical
overhead and no quantum overhead.

## References

  * [1] G. Q. AI, “Suppressing quantum errors by scaling a surface code logical qubit,” _Nature_ , vol. 614, no. 7949, pp. 676–681, 2023.
  * [2] F. Battistel, C. Chamberland, K. Johar, R. W. Overwater, F. Sebastiano, L. Skoric, Y. Ueno, and M. Usman, “Real-time decoding for fault-tolerant quantum computing: Progress, challenges and outlook,” _Nano Futures_ , 2023.
  * [3] A. Bermudez, X. Xu, R. Nigmatullin, J. O’Gorman, V. Negnevitsky, P. Schindler, T. Monz, U. Poschinger, C. Hempel, J. Home _et al._ , “Assessing the progress of trapped-ion processors towards fault-tolerant quantum computation,” _Physical Review X_ , vol. 7, no. 4, p. 041061, 2017.
  * [4] D. Bluvstein, H. Levine, G. Semeghini, T. T. Wang, S. Ebadi, M. Kalinowski, A. Keesling, N. Maskara, H. Pichler, M. Greiner _et al._ , “A quantum processor based on coherent transport of entangled atom arrays,” _Nature_ , vol. 604, no. 7906, pp. 451–456, 2022.
  * [5] S. Boixo, S. V. Isakov, V. N. Smelyanskiy, R. Babbush, N. Ding, Z. Jiang, M. J. Bremner, J. M. Martinis, and H. Neven, “Characterizing quantum supremacy in near-term devices,” _Nature Physics_ , vol. 14, no. 6, pp. 595–600, 2018.
  * [6] N. P. Breuckmann and X. Ni, “Scalable Neural Network Decoders for Higher Dimensional Quantum Codes,” _Quantum_ , vol. 2, p. 68, May 2018. [Online]. Available: https://doi.org/10.22331/q-2018-05-24-68
  * [7] J. J. Burnett, A. Bengtsson, M. Scigliuzzo, D. Niepce, M. Kudra, P. Delsing, and J. Bylander, “Decoherence benchmarking of superconducting qubits,” _npj Quantum Information_ , vol. 5, no. 1, p. 54, 2019.
  * [8] C. C. _et al._ , “Techniques for combining fast local decoders with global decoders under circuit-level noise,” 2022.
  * [9] C. Chamberland, G. Zhu, T. J. Yoder, J. B. Hertzberg, and A. W. Cross, “Topological and subsystem codes on low-degree graphs with flag qubits,” _Physical Review X_ , vol. 10, no. 1, p. 011022, 2020.
  * [10] J. Cheng, H. Wang, Z. Liang, Y. Shi, S. Han, and X. Qian, “Topgen: Topology-aware bottom-up generator for variational quantum circuits,” _arXiv preprint arXiv:2210.08190_ , 2022.
  * [11] I. L. Chuang and M. A. Nielsen, “Prescription for experimental determination of the dynamics of a quantum black box,” _Journal of Modern Optics_ , vol. 44, no. 11-12, pp. 2455–2467, 1997.
  * [12] P. Das, A. Locharla, and C. Jones, “Lilliput: A lightweight low-latency lookup-table decoder for near-term quantum error correction,” in _Proceedings of the 27th ACM International Conference on Architectural Support for Programming Languages and Operating Systems_ , ser. ASPLOS ’22. New York, NY, USA: Association for Computing Machinery, 2022, p. 541–553. [Online]. Available: https://doi.org/10.1145/3503222.3507707
  * [13] P. Das, C. A. Pattison, S. Manne, D. M. Carmean, K. M. Svore, M. Qureshi, and N. Delfosse, “Afs: Accurate, fast, and scalable error-decoding for fault-tolerant quantum computers,” in _2022 IEEE International Symposium on High-Performance Computer Architecture (HPCA)_. IEEE, 2022, pp. 259–273.
  * [14] M. L. Day, P. J. Low, B. White, R. Islam, and C. Senko, “Limits on atomic qubit control from laser noise,” _npj Quantum Information_ , vol. 8, no. 1, p. 72, 2022.
  * [15] N. Delfosse and N. H. Nickerson, “Almost-linear time decoding algorithm for topological codes,” _Quantum_ , vol. 5, p. 595, 2021.
  * [16] N. Delfosse and G. Zémor, “Linear-time maximum likelihood decoding of surface codes over the quantum erasure channel,” _Physical Review Research_ , vol. 2, no. 3, p. 033042, 2020.
  * [17] A. deMarti iOlius, J. E. Martinez, P. Fuentes, and P. M. Crespo, “Performance enhancement of surface codes via recursive minimum-weight perfect-match decoding,” _Physical Review A_ , vol. 108, no. 2, p. 022401, 2023.
  * [18] E. Dennis, A. Kitaev, A. Landahl, and J. Preskill, “Topological quantum memory,” _Journal of Mathematical Physics_ , vol. 43, no. 9, pp. 4452–4505, 2002.
  * [19] J. Edmonds, “Maximum matching and a polyhedron with 0, 1-vertices,” _Journal of research of the National Bureau of Standards B_ , vol. 69, no. 125-130, pp. 55–56, 1965.
  * [20] J. Emerson, R. Alicki, and K. Życzkowski, “Scalable noise estimation with random unitary operators,” _Journal of Optics B: Quantum and Semiclassical Optics_ , vol. 7, no. 10, p. S347, 2005.
  * [21] E. Farhi, J. Goldstone, and S. Gutmann, “A quantum approximate optimization algorithm,” _arXiv preprint arXiv:1411.4028_ , 2014.
  * [22] A. G. Fowler, M. Mariantoni, J. M. Martinis, and A. N. Cleland, “Surface codes: Towards practical large-scale quantum computation,” _Physical Review A_ , vol. 86, no. 3, p. 032324, 2012.
  * [23] A. G. Fowler, A. C. Whiteside, and L. C. Hollenberg, “Towards practical classical processing for the surface code: timing analysis,” _Physical Review A_ , vol. 86, no. 4, p. 042313, 2012.
  * [24] C. Gidney, “Stim: a fast stabilizer circuit simulator,” _Quantum_ , vol. 5, p. 497, 2021.
  * [25] C. Gidney and M. Ekerå, “How to factor 2048 bit rsa integers in 8 hours using 20 million noisy qubits,” _Quantum_ , vol. 5, p. 433, 2021.
  * [26] C. Gidney, M. Newman, A. Fowler, and M. Broughton, “A fault-tolerant honeycomb memory,” _Quantum_ , vol. 5, p. 605, 2021.
  * [27] C. Gidney, M. Newman, and M. McEwen, “Benchmarking the planar honeycomb code,” _Quantum_ , vol. 6, p. 813, 2022.
  * [28] D. Greenbaum, “Introduction to quantum gate set tomography,” _arXiv preprint arXiv:1509.02921_ , 2015.
  * [29] L. K. Grover, “A fast quantum mechanical algorithm for database search,” in _Proceedings of the twenty-eighth annual ACM symposium on Theory of computing_ , 1996, pp. 212–219.
  * [30] E. Gümüş, D. Majidi, D. Nikolić, P. Raif, B. Karimi, J. Peltonen, E. Scheer, J. Pekola, H. Courtois, W. Belzig _et al._ , “Calorimetry of a quantum phase slip,” _arXiv preprint arXiv:2202.08726_ , 2022.
  * [31] S. Gustavsson, F. Yan, G. Catelani, J. Bylander, A. Kamal, J. Birenbaum, D. Hover, D. Rosenberg, G. Samach, A. P. Sears _et al._ , “Suppressing relaxation in superconducting qubits by quasiparticle pumping,” _Science_ , vol. 354, no. 6319, pp. 1573–1577, 2016.
  * [32] J. Haah and M. B. Hastings, “Boundaries for the honeycomb code,” _Quantum_ , vol. 6, p. 693, 2022.
  * [33] M. B. Hastings and J. Haah, “Dynamically generated logical qubits,” _Quantum_ , vol. 5, p. 564, 2021.
  * [34] O. Higgott and C. Gidney, “Sparse blossom: correcting a million errors per core second with minimum-weight matching,” _arXiv preprint arXiv:2303.15933_ , 2023.
  * [35] A. Holmes, M. R. Jokar, G. Pasandi, Y. Ding, M. Pedram, and F. T. Chong, “Nisq+: Boosting quantum computing power by approximating quantum error correction,” in _2020 ACM/IEEE 47th Annual International Symposium on Computer Architecture (ISCA)_. IEEE, 2020, pp. 556–569.
  * [36] A. Kitaev, “Fault-tolerant quantum computation by anyons,” _Annals of Physics_ , vol. 303, no. 1, pp. 2–30, jan 2003. [Online]. Available: https://doi.org/10.1016%2Fs0003-4916%2802%2900018-0
  * [37] I. Kivlichan _et al._ , “Improved fault-tolerant quantum simulation of condensed-phase correlated electrons via trotterization,” _Quantum_ , vol. 4, p. 296, 2020.
  * [38] P. Klimov, J. Kelly, Z. Chen, M. Neeley, A. Megrant, B. Burkett, R. Barends, K. Arya, B. Chiaro, Y. Chen _et al._ , “Fluctuations of energy-relaxation times in superconducting qubits,” _Physical review letters_ , vol. 121, no. 9, p. 090502, 2018.
  * [39] E. Knill, D. Leibfried, R. Reichle, J. Britton, R. B. Blakestad, J. D. Jost, C. Langer, R. Ozeri, S. Seidelin, and D. J. Wineland, “Randomized benchmarking of quantum gates,” _Physical Review A_ , vol. 77, no. 1, p. 012307, 2008.
  * [40] J. Lee _et al._ , “Even more efficient quantum computations of chemistry through tensor hypercontraction,” _PRX Quantum_ , vol. 2, no. 3, p. 030305, 2021.
  * [41] Z. Liang, J. Cheng, H. Ren, H. Wang, F. Hua, Y. Ding, F. Chong, S. Han, Y. Shi, and X. Qian, “Pan: Pulse ansatz on nisq machines,” _arXiv preprint arXiv:2208.01215_ , 2022.
  * [42] Z. Liang, Z. Song, J. Cheng, Z. He, J. Liu, H. Wang, R. Qin, Y. Wang, S. Han, X. Qian, and Y. Shi, “Hybrid gate-pulse model for variational quantum algorithms,” in _2023 60th ACM/IEEE Design Automation Conference (DAC)_ , 2023, pp. 1–6.
  * [43] Z. Liang, Z. Wang, J. Yang, L. Yang, Y. Shi, and W. Jiang, “Can noise on qubits be learned in quantum neural network? a case study on quantumflow,” in _2021 IEEE/ACM International Conference On Computer Aided Design (ICCAD)_. IEEE, 2021, pp. 1–7.
  * [44] S. Lloyd, M. Mohseni, and P. Rebentrost, “Quantum algorithms for supervised and unsupervised machine learning,” _arXiv preprint arXiv:1307.0411_ , 2013.
  * [45] E. Magesan, J. M. Gambetta, and J. Emerson, “Scalable and robust randomized benchmarking of quantum processes,” _Physical review letters_ , vol. 106, no. 18, p. 180504, 2011.
  * [46] M. Mills, J. Sedlacek, T. Peterson, S. Campbell, J. Johansen, J. Dreiling, and D. Francois, “The impact of the sun on trapped-ion quantum computers,” in _2022 IEEE International Conference on Quantum Computing and Engineering (QCE)_ , 2022, pp. 879–881.
  * [47] C. Müller, J. H. Cole, and J. Lisenfeld, “Towards understanding two-level-systems in amorphous solids: insights from quantum circuits,” _Reports on Progress in Physics_ , vol. 82, no. 12, p. 124501, 2019.
  * [48] E. Nielsen, J. K. Gamble, K. Rudinger, T. Scholten, K. Young, and R. Blume-Kohout, “Gate set tomography,” _Quantum_ , vol. 5, p. 557, 2021.
  * [49] M. A. Nielsen and I. Chuang, “Quantum computation and quantum information,” 2002.
  * [50] A. Peruzzo, J. McClean, P. Shadbolt, M.-H. Yung, X.-Q. Zhou, P. J. Love, A. Aspuru-Guzik, and J. L. O’brien, “A variational eigenvalue solver on a photonic quantum processor,” _Nature communications_ , vol. 5, no. 1, pp. 1–7, 2014.
  * [51] J. F. Poyatos, J. I. Cirac, and P. Zoller, “Complete characterization of a quantum process: The two-bit quantum gate,” _Phys. Rev. Lett._ , vol. 78, pp. 390–393, Jan 1997. [Online]. Available: https://link.aps.org/doi/10.1103/PhysRevLett.78.390
  * [52] G. S. Ravi, J. M. Baker, A. Fayyazi, S. F. Lin, A. Javadi-Abhari, M. Pedram, and F. T. Chong, “Have your qec and bandwidth too!: A lightweight cryogenic decoder for common/trivial errors, and efficient bandwidth+ execution management otherwise,” _arXiv preprint arXiv:2208.08547_ , 2022.
  * [53] G. S. Ravi, K. Smith, J. M. Baker, T. Kannan, N. Earnest, A. Javadi-Abhari, H. Hoffmann, and F. T. Chong, “Navigating the dynamic noise landscape of variational quantum algorithms with qismet,” in _Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2_ , 2023, pp. 515–529.
  * [54] C. Ryan-Anderson, J. G. Bohnet, K. Lee, D. Gresh, A. Hankin, J. Gaebler, D. Francois, A. Chernoguzov, D. Lucchetti, N. C. Brown _et al._ , “Realization of real-time fault-tolerant quantum error correction,” _Physical Review X_ , vol. 11, no. 4, p. 041058, 2021.
  * [55] S. Schlör, J. Lisenfeld, C. Müller, A. Bilmes, A. Schneider, D. P. Pappas, A. V. Ustinov, and M. Weides, “Correlating decoherence in transmon qubits: Low frequency noise by single fluctuators,” _Physical review letters_ , vol. 123, no. 19, p. 190502, 2019.
  * [56] P. W. Shor, “Polynomial-time algorithms for prime factorization and discrete logarithms on a quantum computer,” _SIAM review_ , vol. 41, no. 2, pp. 303–332, 1999.
  * [57] Y. Tomita and K. M. Svore, “Low-distance surface codes under realistic quantum noise,” _Physical Review A_ , vol. 90, no. 6, p. 062320, 2014.
  * [58] G. Torlai and R. G. Melko, “Neural decoder for topological codes,” _Phys. Rev. Lett._ , vol. 119, p. 030501, Jul 2017. [Online]. Available: https://link.aps.org/doi/10.1103/PhysRevLett.119.030501
  * [59] S. Vittal, P. Das, and M. Qureshi, “Astrea: Accurate quantum error-decoding via practical minimum-weight perfect-matching,” in _Proceedings of the 50th Annual International Symposium on Computer Architecture_ , ser. ISCA ’23. New York, NY, USA: Association for Computing Machinery, 2023. [Online]. Available: https://doi.org/10.1145/3579371.3589037
  * [60] T. Wagner, H. Kampermann, D. Bruß, and M. Kliesch, “Pauli channels can be estimated from syndrome measurements in quantum error correction,” _Quantum_ , vol. 6, p. 809, 2022.
  * [61] H. Wang, Y. Ding, J. Gu, Y. Lin, D. Z. Pan, F. T. Chong, and S. Han, “Quantumnas: Noise-adaptive search for robust quantum circuits,” in _2022 IEEE International Symposium on High-Performance Computer Architecture (HPCA)_. IEEE, 2022, pp. 692–708.
  * [62] H. Wang, J. Gu, Y. Ding, Z. Li, F. T. Chong, D. Z. Pan, and S. Han, “Quantumnat: quantum noise-aware training with noise injection, quantization and normalization,” in _Proceedings of the 59th ACM/IEEE Design Automation Conference_ , 2022, pp. 1–6.
  * [63] H. Wang, Z. Li, J. Gu, Y. Ding, D. Z. Pan, and S. Han, “Qoc: quantum on-chip training with parameter shift and gradient pruning,” in _Proceedings of the 59th ACM/IEEE Design Automation Conference_ , 2022, pp. 655–660.
  * [64] H. Wang, P. Liu, J. Cheng, Z. Liang, J. Gu, Z. Li, Y. Ding, W. Jiang, Y. Shi, X. Qian _et al._ , “Quest: Graph transformer for quantum circuit reliability estimation,” in _IEEE/ACM International Conference on Computer-Aided Design (ICCAD)_ , 2022.
  * [65] H. Wang, P. Liu, K. Shao, D. Li, J. Gu, D. Z. Pan, Y. Ding, and S. Han, “Transformerqec: Transferable transformer for quantum error correction code decoding,” in _IEEE/ACM International Conference on Computer-Aided Design (ICCAD), FastML for Science Workshop_ , 2023.
  * [66] H. Wang, P. Liu, B. Tan, Y. Liu, J. Gu, D. Z. Pan, J. Cong, U. Acar, and S. Han, “FPQA-C : A Compilation Framework for Field Programmable Qubit Array,” in _arXiv preprint_ , 2023.
  * [67] H. Wang, Y. Liu, P. Liu, J. Gu, Z. Li, Z. Liang, J. Cheng, Y. Ding, X. Qian, Y. Shi, D. Z. Pan, F. T. Chong, and S. Han, “Robuststate: Boosting fidelity of quantum state preparation via noise-aware variational training,” _arXiv preprint_ , 2023.
  * [68] H. Wang, B. Tan, P. Liu, Y. Liu, J. Gu, D. Z. Pan, J. Cong, and S. Han, “Q-Pilot: Field Programmable Quantum Array Compilation with Flying Ancillas,” in _arXiv preprint_ , 2023.
  * [69] Y. Wu, N. Liyanage, and L. Zhong, “An interpretation of union-find decoder on weighted graphs,” _arXiv preprint arXiv:2211.03288_ , 2022.
  * [70] Y. Wu and L. Zhong, “Fusion blossom: Fast mwpm decoders for qec,” _arXiv preprint arXiv:2305.08307_ , 2023.
  * [71] H. Zheng, G. S. Ravi, H. Wang, K. Setia, F. T. Chong, and J. Liu, “Sncqa: A hardware-efficient equivariant quantum convolutional circuit architecture,” in _2023 IEEE International Conference on Quantum Computing and Engineering (QCE)_ , 2023.