# Precedent-Enhanced Legal Judgment Prediction with LLM and Domain-Model
Collaboration

Yiquan Wu1, Siying Zhou1, Yifei Liu1, Weiming Lu1∗, Xiaozhong Liu2
Yating Zhang3, Changlong Sun13, Fei Wu1∗, Kun Kuang1∗
1Zhejiang University, Hangzhou, China
2Worcester Polytechnic Institute, USA
3Alibaba Group, Hangzhou, China
{wuyiquan, zhousiying, liuyifei, luwm<EMAIL_ADDRESS><EMAIL_ADDRESS>
<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Legal Judgment Prediction (LJP) has become an increasingly crucial task in
Legal AI, i.e., predicting the judgment of the case in terms of case fact
description. Precedents are the previous legal cases with similar facts, which
are the basis for the judgment of the subsequent case in national legal
systems. Thus, it is worthwhile to explore the utilization of precedents in
the LJP. Recent advances in deep learning have enabled a variety of techniques
to be used to solve the LJP task. These can be broken down into two
categories: large language models (LLMs) and domain-specific models. LLMs are
capable of interpreting and generating complex natural language, while domain
models are efficient in learning task-specific information. In this paper, we
propose the precedent-enhanced LJP framework (PLJP) – a system that leverages
the strength of both LLM and domain models in the context of precedents.
Specifically, the domain models are designed to provide candidate labels and
find the proper precedents efficiently, and the large models will make the
final prediction with an in-context precedents comprehension. Experiments on
the real-world dataset demonstrate the effectiveness of our PLJP. Moreover,
our work shows a promising direction for LLM and domain-model collaboration
that can be generalized to other vertical domains.

11footnotetext: Corresponding Authors.

##  1 Introduction

Legal AI has been the subject of research for several decades, with the aim of
assisting individuals in various legal tasks, including legal QA (Monroy et
al., 2009), court view generation (Wu et al., 2020), legal entity recognition
(Cardellino et al., 2017), and so on. As one of the most important legal
tasks, legal judgment prediction (LJP) aims to predict the legal judgment of
the case based on the case fact description. The legal judgment typically
includes the law article, charge and prison term.

Figure 1:  An illustration of the judicial process, our motivation is to
promote the collaboration between the domain model and LLM (right part) for
simulating the judicial process of the human judge (left).

Precedents, which refer to previous cases with similar fact descriptions, hold
a crucial position within national legal systems (Guillaume, 2011). On a more
macro level, precedents are known as the collective body of judge-made laws in
a nation(Garner, 2001). They serve the purpose of ensuring consistency in
judicial decisions, providing greater legal guidance to judges and
facilitating legal progress and evolution to meet dynamic legal demands. In
the Common Law system, the precedents are the mandatory basis of the judgment
of the subsequent case (Rigoni, 2014). In the Civil Law system, judge-made
laws are perceived as secondary legal sources while written laws are the basic
legal sources(Larenz, 1992). In the contemporary era, there is also a growing
trend to treat the precedents as a source of “soft” law (Fon and Parisi,
2006), and judges are expected to take them into account when reaching a
decision (Guillaume, 2011). Thus, it is worthwhile to explore the utilization
of precedents in the legal judgment prediction.

With the development of deep learning, many technologies have been adopted in
the LJP task, which can be split into two categories: large language models
(LLMs) and domain-specific models (Ge et al., 2023). Owing to extensive
training, LLMs are good at understanding and generating complex natural
language, as well as in-context learning. On the other hand, domain-specific
models are designed to cater to specific tasks and offer cost-effective
solutions. However, when it comes to incorporating precedents into the LJP
task, both categories of models face certain limitations. LLMs, constrained by
their prompt length, struggle to grasp the meaning of numerous abstract labels
and accurately select the appropriate one. For domain models, though trained
with label annotations, the drawback is the limited ability to comprehend and
distinguish the similarities and differences between the precedents and the
given case.

In this paper, as Fig. 1 shows, we try to collaborate the LLMs with the
domain-specific models and propose a novel precedent-enhanced legal judgment
prediction framework (PLJP). Specifically, domain models contribute by
providing candidate labels and finding the proper precedents from the case
database effectively; the LLMs will decide the final prediction through an in-
context precedent comprehension.

Following the previous LJP works (Zhong et al., 2018; Yue et al., 2021; Dong
and Niu, 2021), our experiments are conducted on the publicly available real-
world legal dataset. To prevent any potential data leakage during the training
of the LLMs, where the model may have already encountered the test cases, we
create a new test set comprising cases that occurred after 2022. This is
necessary because the LLMs we utilize have been trained on a corpus collected
only until September 2021. By doing so, we ensure a fair evaluation of the
PLJP framework. Remarkably, our proposed PLJP framework achieves state-of-the-
art (SOTA) performance on both the original test set and the additional test
set.

To sum up, our main contributions are as follows:

  * •

We address the important task of legal judgment prediction (LJP) by taking
precedents into consideration.

  * •

We propose a novel precedent-enhanced legal judgment prediction (PLJP)
framework that leverages the strength of both LLM and domain models.

  * •

We conduct extensive experiments on the real-world dataset and create an
additional test set to ensure the absence of data leakage during LLM training.
The results obtained on both the original and additional test sets validate
the effectiveness of the PLJP framework.

  * •

Our work shows a promising direction for LLM and domain-model collaboration
that can be generalized over vertical domains. We make all the codes and data
publicly available to motivate other scholars to investigate this novel and
interesting research direction111The github link is hidden for the anonymous
review..

##  2 Related Work

###  2.1 Legal AI

Legal Artificial Intelligence (Legal AI) aims to enhance tasks within the
legal domain through the utilization of artificial intelligence techniques
(Zhong et al., 2020; Katz et al., 2023). Collaborative efforts between
researchers in both law and computer fields have been lasting to explore the
potential of Legal AI and its applications across various legal tasks. These
tasks encompass areas such as legal question answering (QA) (Monroy et al.,
2009), legal entity recognition (Cardellino et al., 2017), court view
generation (Wu et al., 2020), legal summarization (Hachey and Grover, 2006;
Bhattacharya et al., 2019), legal language understanding(Chalkidis et al.,
2022) and so on.

In this work, we focus on the task of legal judgment prediction, which is one
of the most common tasks in Legal AI.

###  2.2 Legal Judgment Prediction

Legal judgment prediction (LJP) aims to predict judgment results based on the
fact descriptions automatically (Lin et al., 2012; Chalkidis et al., 2019; Yue
et al., 2021; Xu et al., 2020; Niklaus et al., 2021; Malik et al., 2021; Feng
et al., 2022; Lyu et al., 2022; Gan et al., 2022). The LJP methods in earlier
years required manually extracted features (Keown, 1980), which is simple but
costly. Owing to the prosperity of machine learning (Wu et al., 2022; Shen et
al., 2022; Li et al., 2022a, b; Zhang et al., 2022; Li et al., 2023; Zhang et
al., 2023), researchers began to formalize the LJP problem with machine
learning methods. These data-driven methods can learn the features with far
less labor (e.g., only the final labels are required). Sulea et al. (2017)
developed an ensemble system that averages the output of multiple SVM to
improve the performance of LJP. Luo et al. (2017) utilized an attention
mechanism in the LJP. Zhong et al. (2018) considered the dependency of the
sub-tasks in the LJP. Yue et al. (2021) investigated the problem by separating
the representation of fact description into different embedding. Liu et al.
(2022) used contrastive learning in the LJP.

However, these existing LJP methods tend to overlook the significance of
precedents. In this study, we propose a precedent-enhanced LJP framework
(PLJP) that leverages the collaboration between domain-specific models and
large language models (LLMs) to address the LJP task.

###  2.3 Precedent Retrieval

The precedent is the basis of judgment in the Common Law system, and also an
important reference for decision-making in the Civil Law system. Therefore,
precedent retrieval is another valuable task in Legal AI (Althammer et al.,
2021). There are two main precedent retrieval models: expert knowledge-based
models and natural language processing (NLP)-based models (Bench-Capon et al.,
2012). Expert knowledge-based models use the designed sub-elements to
represent the legal cases (Saravanan et al., 2009), while NLP-based models
mainly convert the text into embeddings and then calculate the similarity from
the embedding level (Ma et al., 2021; Chalkidis et al., 2020).

Most retrieval models required additional annotation so can not be directly
applied to the LJP task. In our paper, we use an unsupervised dense retrieval
model (Izacard et al., 2022) to get the precedents, which can be updated by
other retrieval models if needed.

###  2.4 Large Language Models

Large language models (LLMs), such as ChatGPT, have attracted widespread
attention from society (Zhao et al., 2023). With pre-training over large-scale
corpora, LLMs show strong capabilities in interpreting and generating complex
natural language, as well as reasoning (e.g., in-context learning). The
technical evolution of LLMs has been making an important impact on the fields
of natural language processing (Brown et al., 2020; Touvron et al., 2023),
computer vision (Shao et al., 2023; Wu et al., 2023), and reinforcement
learning (Du et al., 2023). In the legal domain, LLMs can also be used for
many tasks such as legal document analysis and legal document writing (Sun,
2023).

However, in the prediction tasks, which can involve dozens of abstract labels,
the performance of LLMs is not as good as in generation tasks, due to the
limited prompt length. In this paper, we explore the utilization of LLMs in
the LJP task with the collaboration of domain-specific models.

##  3 Problem Formulation

In this work, we focus on the problem of legal judgment prediction. We first
clarify the definition of the terms as follows.

$\bullet$ Fact Description refers to a concise narrative of the case, which
typically includes the timeline of events, the actions or conduct of each
party, and any other essential details that are relevant to the case. Here we
define it as a token sequence $f=\\{w_{t}^{f}\\}_{t=1}^{l_{f}}$, where $l_{f}$
is the length.

$\bullet$ Judgment is the final decision made by a judge in a legal case based
on the facts and the precedents. It typically consists of the law article, the
charge, and the prison term. We represent the judgment of a case as
$j=(a,c,t)$, where $a$, $c$, $t$ refer to the labels of article, charge and
prison term, respectively.

$\bullet$ Precedent is the previous case with a similar fact. The judgments of
the precedents are important references for the current case. Here, a
precedent is defined as $p=(f_{p},j_{p})$, where $f_{p}$ is its fact
description and $j_{p}$ is its judgment. For a given case, there can be
several precedents, which can be denoted as $P=\\{p_{1},p_{2},...,p_{n}\\}$,
where $n$ is the number of precedents.

Then the problem can be defined as:

######  Problem 1 (Legal Judgment Prediction).

Given the fact description $f$, our task is to get and comprehend the
precedents $P$, then predict the judgment $j=(a,c,t)$.

##  4 Precedent-Enhanced LJP (PLJP)

In this section, we describe our precedent-enhanced legal judgment prediction
framework (PLJP), Fig. 2 shows the overall framework.

Figure 2: The overall framework of PLJP, where the sub, obj and ex refer to
the subjective motivation, objective behavior and ex post facto circumstance,
respectively. The solid lines are the precedent retrieval process, while the
dotted lines represent the process of the prediction.

###  4.1 Case Database Construction

Before we use the precedents, we have to collect a large number of previous
cases to construct a case database. Since the fact descriptions are usually
long and elaborate, it is difficult for the models to get the proper
precedents. To this end, we reorganize the fact description of these previous
cases with the help of LLMs.

####  4.1.1 Fact Reorganization

Given a fact description of a case, we summarize it from three aspects:
subjective motivation, objective behavior, and ex post facto circumstances.
The reorganization doesn’t require human annotation and is completed by the
LLMs with the following prompts:  “A fact description can be categorized into
subjective motivation, objective behavior, and ex post facto circumstances.
Subjective motivation refers to the psychological attitude of the perpetrator
towards their harmful actions and their consequences, including intent,
negligence, and purposes of the crime. Objective behavior pertains to the
necessary conditions for constituting a crime in terms of observable
activities, including harmful conduct, harmful results, and the causal
relationship between the conduct and the results. Ex post facto circumstances
are various factual situations considered when determining the severity of
penalties. Mitigating circumstances for lenient punishment include voluntary
surrender and meritorious conduct, while aggravating circumstances for harsher
punishment include recidivism. Based on the provided information, your task is
to summarize the following facts.”

The reorganization reduces the length of facts and makes the precedents easy
to get and comprehend in the PLJP.

After the reorganization, the fact description $f$ is translated to a triplet
$(sub,obj,ex)$, which indicates the subjective motivation, objective behavior,
and ex post facto circumstances, respectively. Finally, a previous case in the
case database is stored as a pair of reorganized facts and the judgment.

###  4.2 Legal Judgment Prediction

Next, we describe the collaboration of the LLM and domain models in legal
judgment prediction.

####  4.2.1 Domain Models

The domain models are trained on specific datasets, aiming to solve certain
tasks. Here, we use two kinds of domain models, including the predictive model
and the retrieval model.

##### Predictive model.

The predictive model takes the fact description as the input and outputs the
candidate labels of the three sub-tasks (e.g., law article, charge, prison
term). Since the fact description $f=\\{w_{t}^{f}\\}_{t=1}^{l_{f}}$ are
sequences of words, we first transform it into embedding sequence
$H^{f}\in\mathbb{R}^{l_{f}\times d}$ with an Encoder:

$H^{f}=\operatorname{Encode}(f),$ (1)

where $H^{f}={h^{f}_{1},h^{f}_{2},...,h^{f}_{l_{f}}}$, and $d$ is the
dimension of the embedding.

We take a max-pooling operation to obtain the pooled hidden vector
$h^{f}\in\mathbb{R}^{d}$ and then feed it into a fully-connected network with
softmax activation to obtain the label probability distribution
$P\in\mathbb{R}^{m}$:

$\displaystyle h^{f}$ $\displaystyle=\operatorname{MaxPooling}(H^{f}),$ (2)
$\displaystyle P$ $\displaystyle=\operatorname{Softmax}(W^{p}\cdot
h^{f}+b^{p}),$

where $W^{p}\in\mathbb{R}^{m\times d}$ and $b^{p}\in\mathbb{R}^{m}$ are
learnable parameters. Note $m$ varies in different sub-tasks.

Then, each sub-task gets its candidate labels according to the probability
distribution $P$, and the number of candidate labels is equal to the number of
precedents $n$.

##### Retrieval model.

The retrieval model aims to get the proper precedents of the given case based
on its reorganized fact $(sub,obj,ex)$.

Formally, to get the similarity score of any two texts $D_{1}$ and $D_{2}$, we
will first encode each of them independently using the same encoder:

$h_{D_{1}}=\operatorname{Encoder}(D_{1}),h_{D_{2}}=\operatorname{Encoder}(D_{2}),$
(3)

where $h_{D_{1}}\in\mathbb{R}^{d^{\prime}}$ and
$h_{D_{2}}\in\mathbb{R}^{d^{\prime}}$ are the embedding of each, $d^{\prime}$
is the dimension. The similarity score $s(D_{1},D_{2})$ is then the cosine
similarity of the $h_{D_{1}}$ and $h_{D_{2}}$:

$s(D_{1},D_{2})=\frac{h_{D_{1}}\cdot
h_{D_{2}}}{\left\|h_{D_{1}}\right\|\left\|h_{D_{2}}\right\|}.$ (4)

Here we concatenate the $sub$, $obj$ and $ex$ into a whole text to calculate
the similarity score of the given case and the cases in the case database.

For each candidate label, we pick one case as the precedent: the case that has
the highest similarity score and has the same label. For example, if the label
“Theft” is in the candidate labels in the charge prediction, we will find the
most similar previous case with the same label as the corresponding precedent.
The one-to-one relationship between the candidate label and precedent helps
the LLM distinguish the differences among the labels. In other words, the
precedent serves as a supplementary explanation of the label.

Finally, we get precedents $P=\\{p_{1},p_{2},...,p_{n}\\}$ for the given case.

####  4.2.2 LLMs

The large language models are models with billions of parameters, which are
trained on large-scale corpora, and show strong capabilities in interpreting
and generating complex natural language. LLMs contribute to PLJP by fact
reorganization and in-context precedent comprehension.

##### Fact Reorganization

The fact reorganization is described in case database construction (Sec.
4.1.1), which aims to summarize the fact description from three aspects by the
LLMs. Besides the database contribution, as Fig. 2 shows, when a new test case
comes, the LLMs will reorganize the fact description with the same prompt.

##### In-Context Precedent Comprehension

Since LLMs are capable of understanding complex natural language, we stack the
given case with its precedents and let the LLMs make the final prediction by
an in-context precedent comprehension. Specifically, the prompt of law article
prediction is designed as follows:  “Based on the facts, we select the
candidate law articles by the domain models and select the following three
precedents based on the candidate law articles. Please comprehend the
difference among the precedents, then compare them with the facts of this
case, and choose the final label.”

Consider the topological dependencies among the three sub-tasks (Zhong et al.,
2018), in the prediction of charge, we add the predicted law article in the
prompt; in the prediction of prison term, we add the predicted law article and
charge.

###  4.3 Training

In PLJP, considering the realizability, we train domain models on legal
datasets and leave the LLMs unchanged. To train predictive models, the cross-
entropy loss is employed. As for retrieval models, contrastive loss is used
like Izacard et al. (2022).

##  5 Experiments

Type | CAIL2018 | CJO22
---|---|---
# Law Article | 164 | 164
# Charge | 42 | 42
# Prison Term | 10 | 10
# Sample | 82138 | 1698
Avg. # words in Fact | 288.6 | 461.7
Table 1: Statistics of datasets. Method | CJO22 | CAIL2018
---|---|---
Acc | Ma-P | Ma-R | Ma-F | Acc | Ma-P | Ma-R | Ma-F
CNN (LeCun et al., 1989) | 76.14 | 35.48 | 38.55 | 35.39 | 80.50 | 40.10 | 38.33 | 38.49
BERT (Devlin et al., 2019) | 82.62 | 45.89 | 47.91 | 45.83 | 82.77 | 36.82 | 35.94 | 35.82
Roberta (Liu et al., 2019) | 80.32 | 42.36 | 44.22 | 41.80 | 83.08 | 48.09 | 44.25 | 44.87
TopJudge (Zhong et al., 2018) | 78.73 | 40.38 | 41.47 | 40.09 | 80.46 | 40.96 | 40.96 | 38.24
R-Former (Dong and Niu, 2021) | 87.69 | 53.03 | 49.35 | 50.23 | 87.82 | 56.13 | 56.57 | 55.81
LADAN (Xu et al., 2020) | 79.44 | 48.43 | 44.13 | 46.18 | 82.82 | 42.57 | 39.00 | 40.71
NeurJudge (Yue et al., 2021) | 71.38 | 52.86 | 53.52 | 52.62 | 76.91 | 55.95 | 52.92 | 53.56
EPM(Feng et al., 2022) | 84.19 | 47.21 | 43.79 | 44.39 | 85.80 | 49.08 | 45.76 | 47.32
CTM(Liu et al., 2022) | 79.44 | 47.83 | 42.25 | 43.43 | 84.72 | 46.46 | 44.83 | 45.10
Dav003 | 2.10 | 0.82 | 0.17 | 0.26 | 1.02 | 0.30 | 0.08 | 0.13
3.5turbo | 9.13 | 2.54 | 1.61 | 1.53 | 4.08 | 4.95 | 3.64 | 2.30
PLJP(CNN) | 87.67 | 55.21 | 55.59 | 54.37 | 86.05 | 58.08 | 56.46 | 54.92
PLJP(BERT) | 94.18 | 74.65 | 76.23 | 74.84 | 87.07 | 58.81 | 57.29 | 56.63
Table 2: Results of law article prediction, the best is bolded and the second best is underlined. Method | CJO22 | CAIL2018
---|---|---
Acc | Ma-P | Ma-R | Ma-F | Acc | Ma-P | Ma-R | Ma-F
CNN (LeCun et al., 1989) | 74.91 | 74.00 | 78.12 | 73.97 | 87.52 | 88.23 | 88.31 | 88.17
BERT (Devlin et al., 2019) | 80.50 | 80.34 | 81.09 | 78.36 | 89.10 | 90.10 | 89.48 | 89.63
Roberta (Liu et al., 2019) | 79.26 | 78.93 | 81.25 | 78.18 | 90.30 | 91.02 | 90.97 | 90.94
TopJudge (Zhong et al., 2018) | 76.67 | 74.00 | 77.40 | 74.62 | 87.31 | 88.68 | 87.84 | 88.20
R-Former (Dong and Niu, 2021) | 90.71 | 93.06 | 88.66 | 89.82 | 91.54 | 91.61 | 91.96 | 91.58
LADAN (Xu et al., 2020) | 79.64 | 48.43 | 44.13 | 46.18 | 88.09 | 90.12 | 88.82 | 89.47
NeurJudge (Yue et al., 2021) | 71.85 | 69.37 | 71.09 | 68.66 | 82.13 | 82.71 | 82.30 | 82.36
EPM(Feng et al., 2022) | 83.49 | 80.36 | 83.29 | 81.87 | 91.20 | 90.81 | 89.99 | 90.46
CTM(Liu et al., 2022) | 79.33 | 82.39 | 83.12 | 82.81 | 90.28 | 90.34 | 88.08 | 86.30
Dav003 | 44.65 | 52.43 | 32.93 | 35.29 | 25.85 | 35.37 | 25.09 | 22.08
3.5turbo | 58.37 | 56.03 | 40.68 | 42.62 | 49.65 | 42.29 | 34.05 | 31.85
PLJP(CNN) | 91.62 | 83.43 | 84.88 | 83.40 | 91.49 | 81.80 | 83.95 | 80.06
PLJP(BERT) | 94.18 | 90.25 | 88.67 | 89.05 | 94.99 | 92.12 | 91.10 | 91.33
Table 3: Results of charge prediction, the best is bolded and the second best
is underlined.

###  5.1 Datasets

Following many influential LJP works (Zhong et al., 2018; Xu et al., 2020; Yue
et al., 2021; Dong and Niu, 2021), our experiment is conducted on the widely
used and publicly available CAIL2018 dataset, which is a Chinese dataset in
the context of People’s Republic of China (PRC). This dataset consists of
real-world cases, each of which includes a fact description accompanied by a
complete judgment encompassing three labels: law articles, charges, and prison
terms222Prison terms are divided into non-overlapping intervals..

To mitigate the potential data leakage during the training of LLMs, which were
trained on corpora collected until September 2021, we have compiled a new
dataset called CJO22. This dataset exclusively contains legal cases that
occurred after 2022, sourced from the same origin as
CAIL2018333https://wenshu.court.gov.cn/. However, due to its limited size, the
newly collected CJO22 dataset is inadequate for the training purposes of the
domain models. Consequently, we utilize it solely as an additional test set.
To facilitate meaningful comparisons, we retain only the labels that are
common to both datasets, considering that the labels may not be entirely
aligned.

Tab. 1 shows the statistics of the processed datasets, and all the experiments
are conducted on the same datasets. For CAIL2018 dataset, we randomly divide
it into training set, validation set and test set according to the ratio of 8:
1: 1.

The previous cases in the case database are sampled from the training dataset,
and we set the amount to 4000.

Method | CJO22 | CAIL2018
---|---|---
Acc | Ma-P | Ma-R | Ma-F | Acc | Ma-P | Ma-R | Ma-F
CNN (LeCun et al., 1989) | 27.38 | 18.48 | 17.51 | 17.44 | 34.42 | 32.22 | 30.53 | 31.05
BERT (Devlin et al., 2019) | 36.80 | 29.83 | 27.50 | 27.03 | 40.00 | 37.53 | 33.66 | 33.58
Roberta (Liu et al., 2019) | 29.74 | 24.73 | 24.76 | 23.22 | 40.84 | 38.62 | 38.55 | 38.50
TopJudge (Zhong et al., 2018) | 27.14 | 19.76 | 17.69 | 17.94 | 35.54 | 33.55 | 31.08 | 32.00
R-Former (Dong and Niu, 2021) | 38.63 | 32.63 | 32.76 | 29.51 | 40.70 | 36.09 | 36.76 | 35.04
LADAN (Xu et al., 2020) | 33.69 | 26.40 | 22.94 | 24.55 | 38.03 | 33.66 | 30.08 | 31.77
NeurJudge (Yue et al., 2021) | 26.80 | 26.81 | 26.85 | 25.97 | 33.53 | 36.46 | 37.26 | 36.53
EPM(Feng et al., 2022) | 36.91 | 30.65 | 31.61 | 30.20 | 40.25 | 37.96 | 37.00 | 37.34
CTM(Liu et al., 2022) | 36.81 | 27.10 | 25.96 | 26.46 | 39.56 | 38.66 | 38.02 | 37.84
Dav003 | 0.47 | 5.56 | 0.21 | 0.41 | 0.68 | 10.38 | 0.49 | 0.94
3.5turbo | 1.40 | 1.16 | 1.07 | 1.11 | 1.02 | 2.71 | 1.13 | 1.15
PLJP(CNN) | 36.51 | 20.21 | 21.44 | 20.07 | 40.81 | 32.77 | 35.59 | 25.71
PLJP(BERT) | 43.52 | 33.37 | 35.67 | 31.98 | 48.72 | 42.64 | 36.80 | 35.43
Table 4: Results of prison term prediction, the best is bolded and the second best is underlined. Method | CJO22 | CAIL2018
---|---|---
Law Article | Charge | Prison Term | Law Article | Charge | Prison Term
Acc | Ma-F | Acc | Ma-F | Acc | Ma-F | Acc | Ma-F | Acc | Ma-F | Acc | Ma-F
w/o p | 54.65 | 28.32 | 83.48 | 76.33 | 35.81 | 20.84 | 85.03 | 51.54 | 85.03 | 70.07 | 32.31 | 22.58
w/o c | 45.34 | 40.22 | 42.32 | 41.85 | 32.55 | 20.26 | 67.35 | 46.65 | 72.79 | 60.34 | 26.53 | 13.66
w/o d | 94.18 | 74.84 | 85.58 | 70.50 | 39.53 | 20.31 | 87.07 | 56.63 | 87.41 | 73.45 | 38.09 | 21.44
w/o r | 88.13 | 58.75 | 87.67 | 74.83 | 36.27 | 23.70 | 86.05 | 58.26 | 86.73 | 77.53 | 38.10 | 21.70
w/ e | 90.70 | 67.90 | 80.70 | 66.53 | 35.35 | 20.21 | 89.80 | 61.64 | 85.37 | 68.48 | 38.44 | 23.14
PLJP | 94.18 | 74.84 | 94.18 | 89.05 | 43.52 | 31.98 | 87.07 | 56.63 | 94.99 | 91.33 | 48.72 | 35.43
Table 5: Results of ablation experiments, the best is bolded and the second
best is underlined. Figure 3: The Ma-F of PLJP with different number of
precedents.

###  5.2 Baselines

For domain-specific LJP baselines, we implement the following for comparison:

CNN (LeCun et al., 1989) extracts text features through convolutional
operations with different kernels for text classification; BERT(Devlin et al.,
2019) is a pre-trained language model and can be easily fine-tuned on the
downstream tasks; TopJudge (Zhong et al., 2018) use multi-task learning and
capture the dependencies among the three sub-task in LJP; NeurJudge (Yue et
al., 2021) splits the fact description into different parts for making
predictions; R-Former (Dong and Niu, 2021) formalizes LJP as a node
classification problem over a global consistency graph and relational learning
is introduced; LADAN (Xu et al., 2020) uses graph distillation to extract
discriminative features of the fact Retri-BERT (Chalkidis and
Kementchedjhieva, 2023) retrieves similar documents to augment the input
document representation for multi-label text classification; EPM (Feng et al.,
2022) locates event-related information essential for judgment while utilizing
cross-task consistency constraints among the subtasks; CTM (Liu et al., 2022)
establishes a LJP framework with case triple modeling from contrastive case
relations.

We use the LLM baselines as follows444We give a fixed example in the prompt to
help the LLMs understand the tasks.: Dav003 means the text-davinci-003,
3.5turbo means the gpt-3.5-turbo. These LLMs are both from the GPT-3.5 family,
released by OpenAI and can understand and generate complex natural
language555https://platform.openai.com/docs/models.

Figure 4: The charge prediction of a given case. The green parts are useful
information for prediction, while the red parts are content that can be
confused by the domain models.

For PLJP, we take the CNN and BERT as the predictive models, and take the
text-davinci-003 as the implementation of the LLM, named as PLJP(CNN) and
PLJP(BERT). The top-k accuracy of CNN and BERT is shown in the Appendix.
Considering the length limit of the prompt, we set the number of precedents to
3.

We also do ablation experiments as follows: PLJP w/o p refers to the removal
of precedents, and the prediction of labels is done solely based on the
candidate labels using the LLM; PLJP w/o c denotes we remove the candidate
labels and predict the label only with the fact description and precedents;
PLJP w/o d means we predict the three labels independently instead of
considering the dependencies among the three subtasks; PLJP w/o r denotes we
find precedents based the raw fact instead of from the reorganized fact; PLJP
w/ e means we let the LLMs generate the explanation of the prediction as well.

In the ablation study, PLJP means PLJP(BERT).

###  5.3 Experiment Settings

Here we describe the implementation of PLJP in our experiments. Note all the
LLMs and domain models are replaceable in the PLJP framework.

In the experiments, for the LLMs, we directly use the APIs provided by OpenAI.
For the domain models, we use the unsupervised dense retrieval model (Izacard
et al., 2022) in precedent retrieval, which gets the precedents from the case
database according to the reorganized facts. For other domain models such as
TopJudge and NeurJudge, we use the training settings from the original paper.

For the metrics, we employ Accuracy (Acc), Macro-Precision (Ma-P), Macro-
Recall (Ma-R) and Macro-F1 (Ma-F).

###  5.4 Experiment Results

We analyze the experimental results in this section.

##### Result of judgment prediction:

From Tab. 2, Tab. 3 and Tab. 4, we have the following observations: 1) The
LLMs perform not well in the prediction tasks alone, especially when the label
has no actual meaning (e.g., the index of the law article and prison term). 2)
By applying our PLJP framework with the collaboration of LLMs and domain
models, the simple models (e.g., CNN, BERT) gain significant improvement. 3)
The model performance on CJO22 is lower than that on CAIL2018, which shows the
challenge of the newly constructed test set. 4) PLJP(BERT) achieves the best
performance in almost all the metric evaluation metrics in both CAIL2018 and
CJO22 test sets, which proves the effectiveness of the PLJP. 5) Compared to
the prediction of the law article and charge, the prediction of prison term is
still a more challenging task. 6) The reported results of the LJP baselines
are not as good as the original papers, this may be because we keep all the
low-frequency labels instead of removing them as the original papers did.

##### Results of ablation experiment:

From Tab. 5, we can conclude that: 1) The performance gap of the PLJP w/o p
and PLJP demonstrates the effects of the precedents. 2) The results of PLJP
w/o c prove the importance of the candidate labels. 3) Considering the
topological dependence of the three sub-tasks benefits the model performance
as PLJP w/o d shows. 4) When we use the raw fact instead of the reorganized
fact, the performance drops (e.g., the Acc of prison term in CJO22 drops from
45.32% to 36.27%). 5) If we force the LLMs to generate the explanation of the
prediction, the performance also drops a bit. We put cases with explanations
in the Appendix.

From Fig. 3, we can find that the performance of PLJP improves as the number
of precedents increases, which also proves the effectiveness of injecting
precedents into the LJP.

###  5.5 Case Study

Fig. 4 shows an intuitive comparison among the three methods in the process of
charge prediction. Based on the fact description of the given case, the domain
models provide candidate charges with the corresponding precedents. As the
case shows, the defendant made fraud by selling the cars that were rented from
other people. However, since there contains “contract” in the fact
description, baselines (e.g., R-Former and BERT) can be misled and predict the
wrong charge of “Contract Fraud”. Through an in-context precedent
comprehension by the LLMs, PLJP(BERT) distinguishes the differences among the
precedents and the given case (e.g., the crime does not occur during the
contracting process, and the contract is only a means to commit the crime),
and give the right result of “Fraud”.

##  6 Conclusion and Future Work

In this paper, we address the important task of legal judgment prediction
(LJP) by taking precedents into consideration. We propose a novel framework
called precedent-enhanced legal judgment prediction (PLJP), which combines the
strength of both LLMs and domain models to better utilize (e.g., retrieve and
comprehend) the precedents. Experiments on the real-world dataset prove the
effectiveness of the PLJP.

Based on the PLJP, in the future, we can explore the following directions: 1)
Develop methods to identify and mitigate any biases that could affect the
predictions and ensure fair and equitable outcomes. 2) Validate the
effectiveness of LLM and domain collaboration in other vertical domains such
as medicine and education.

###  6.1 Ethical Discussion

With the increasing adoption of Legal AI in the field of legal justice, there
has been a growing awareness of the ethical implications involved. The
potential for even minor errors or biases in AI-powered systems can lead to
significant consequences.

In light of these concerns, we have to claim that our work is an algorithmic
exploration and will not be directly used in court so far. Our goal is to
provide suggestions to judges rather than making final judgments without human
intervention. In practical use, human judges should be the final safeguard to
protect justice fairness. In the future, we plan to study how to identify and
mitigate potential biases to ensure the fairness of the model.

##  7 Limitations

In this section, we discuss the limitations of our works as follow:

$\bullet$ We only interact with the LLMs one round per time. The LLMs are
capable of multi-round interaction (e.g., Though of Chains), which may help
the LLM to better understand the LJP task.

$\bullet$ We validate the effectiveness of LLM and domain model collaboration
in the legal domain. It’s worthwhile to explore such collaboration in other
vertical domains such as medicine and education, as well as in other legal
datasets (e.g., the datasets from the Common Law system).

## References

  * Althammer et al. (2021) Sophia Althammer, Arian Askari, Suzan Verberne, and Allan Hanbury. 2021.  Dossier@coliee 2021: Leveraging dense retrieval and summarization-based re-ranking for case law retrieval.  _CoRR_ , abs/2108.03937.
  * Bench-Capon et al. (2012) Trevor J. M. Bench-Capon, Michal Araszkiewicz, Kevin D. Ashley, Katie Atkinson, Floris Bex, Filipe Borges, Danièle Bourcier, Paul Bourgine, Jack G. Conrad, Enrico Francesconi, Thomas F. Gordon, Guido Governatori, Jochen L. Leidner, David D. Lewis, Ronald Prescott Loui, L. Thorne McCarty, Henry Prakken, Frank Schilder, Erich Schweighofer, Paul Thompson, Alex Tyrrell, Bart Verheij, Douglas N. Walton, and Adam Z. Wyner. 2012.  A history of AI and law in 50 papers: 25 years of the international conference on AI and law.  _Artif. Intell. Law_ , 20(3):215–319.
  * Bhattacharya et al. (2019) Paheli Bhattacharya, Kaustubh Hiware, Subham Rajgaria, Nilay Pochhi, Kripabandhu Ghosh, and Saptarshi Ghosh. 2019.  A comparative study of summarization algorithms applied to legal case judgments.  In _Advances in Information Retrieval - 41st European Conference on IR Research, ECIR 2019, Cologne, Germany, April 14-18, 2019, Proceedings, Part I_ , volume 11437 of _Lecture Notes in Computer Science_ , pages 413–428. Springer.
  * Brown et al. (2020) Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. 2020.  Language models are few-shot learners.  In _Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual_.
  * Cardellino et al. (2017) Cristian Cardellino, Milagro Teruel, Laura Alonso Alemany, and Serena Villata. 2017.  Legal NERC with ontologies, wikipedia and curriculum learning.  In _Proceedings of the 15th Conference of the European Chapter of the Association for Computational Linguistics, EACL 2017, Valencia, Spain, April 3-7, 2017, Volume 2: Short Papers_ , pages 254–259. Association for Computational Linguistics.
  * Chalkidis et al. (2019) Ilias Chalkidis, Ion Androutsopoulos, and Nikolaos Aletras. 2019.  Neural legal judgment prediction in english.  In _Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July 28- August 2, 2019, Volume 1: Long Papers_ , pages 4317–4323. Association for Computational Linguistics.
  * Chalkidis et al. (2020) Ilias Chalkidis, Manos Fergadiotis, Prodromos Malakasiotis, Nikolaos Aletras, and Ion Androutsopoulos. 2020.  LEGAL-BERT: the muppets straight out of law school.  _CoRR_ , abs/2010.02559.
  * Chalkidis et al. (2022) Ilias Chalkidis, Abhik Jana, Dirk Hartung, Michael J. Bommarito II, Ion Androutsopoulos, Daniel Martin Katz, and Nikolaos Aletras. 2022.  Lexglue: A benchmark dataset for legal language understanding in english.  In _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2022, Dublin, Ireland, May 22-27, 2022_ , pages 4310–4330. Association for Computational Linguistics.
  * Chalkidis and Kementchedjhieva (2023) Ilias Chalkidis and Yova Kementchedjhieva. 2023.  Retrieval-augmented multi-label text classification.  _CoRR_ , abs/2305.13058.
  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019.  BERT: pre-training of deep bidirectional transformers for language understanding.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers)_ , pages 4171–4186. Association for Computational Linguistics.
  * Dong and Niu (2021) Qian Dong and Shuzi Niu. 2021.  Legal judgment prediction via relational learning.  In _SIGIR ’21: The 44th International ACM SIGIR Conference on Research and Development in Information Retrieval, Virtual Event, Canada, July 11-15, 2021_ , pages 983–992. ACM.
  * Du et al. (2023) Yuqing Du, Olivia Watkins, Zihan Wang, Cédric Colas, Trevor Darrell, Pieter Abbeel, Abhishek Gupta, and Jacob Andreas. 2023.  Guiding pretraining in reinforcement learning with large language models.  _CoRR_ , abs/2302.06692.
  * Feng et al. (2022) Yi Feng, Chuanyi Li, and Vincent Ng. 2022.  Legal judgment prediction: A survey of the state of the art.  In _Proceedings of the Thirty-First International Joint Conference on Artificial Intelligence, IJCAI 2022, Vienna, Austria, 23-29 July 2022_ , pages 5461–5469. ijcai.org.
  * Fon and Parisi (2006) Vincy Fon and Francesco Parisi. 2006.  Judicial precedents in civil law systems: A dynamic analysis.  _International Review of Law and Economics_ , 26(4):519–535.
  * Gan et al. (2022) Leilei Gan, Baokui Li, Kun Kuang, Yi Yang, and Fei Wu. 2022.  Exploiting contrastive learning and numerical evidence for improving confusing legal judgment prediction.  _CoRR_ , abs/2211.08238.
  * Garner (2001) Bryan A. Garner. 2001.  _A dictionary of modern legal usage_.  Oxford University Press, New York.
  * Ge et al. (2023) Yingqiang Ge, Wenyue Hua, Jianchao Ji, Juntao Tan, Shuyuan Xu, and Yongfeng Zhang. 2023.  Openagi: When LLM meets domain experts.  _CoRR_ , abs/2304.04370.
  * Guillaume (2011) Gilbert Guillaume. 2011.  The use of precedent by international judges and arbitrators.  _Journal of International Dispute Settlement_ , 2(1):5–23.
  * Hachey and Grover (2006) Ben Hachey and Claire Grover. 2006.  Extractive summarisation of legal texts.  _Artif. Intell. Law_ , 14(4):305–345.
  * Izacard et al. (2022) Gautier Izacard, Mathilde Caron, Lucas Hosseini, Sebastian Riedel, Piotr Bojanowski, Armand Joulin, and Edouard Grave. 2022.  Unsupervised dense information retrieval with contrastive learning.  _Trans. Mach. Learn. Res._ , 2022.
  * Katz et al. (2023) Daniel Martin Katz, Dirk Hartung, Lauritz Gerlach, Abhik Jana, and Michael J Bommarito II. 2023.  Natural language processing in the legal domain.  _arXiv preprint arXiv:2302.12039_.
  * Keown (1980) R Keown. 1980.  Mathematical models for legal prediction.  _Computer/lj_ , 2:829.
  * Larenz (1992) Karl Larenz. 1992.  _Methodenlehre der Rechtswissenschaft_.  Springer, Berlin, Heidelberg.
  * LeCun et al. (1989) Yann LeCun, Bernhard E. Boser, John S. Denker, Donnie Henderson, Richard E. Howard, Wayne E. Hubbard, and Lawrence D. Jackel. 1989.  Backpropagation applied to handwritten zip code recognition.  _Neural Comput._ , 1(4):541–551.
  * Li et al. (2023) Mengze Li, Han Wang, Wenqiao Zhang, Jiaxu Miao, Zhou Zhao, Shengyu Zhang, Wei Ji, and Fei Wu. 2023.  WINNER: weakly-supervised hierarchical decomposition and alignment for spatio-temporal video grounding.  In _IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2023, Vancouver, BC, Canada, June 17-24, 2023_ , pages 23090–23099. IEEE.
  * Li et al. (2022a) Mengze Li, Tianbao Wang, Haoyu Zhang, Shengyu Zhang, Zhou Zhao, Jiaxu Miao, Wenqiao Zhang, Wenming Tan, Jin Wang, Peng Wang, Shiliang Pu, and Fei Wu. 2022a.  End-to-end modeling via information tree for one-shot natural language spatial video grounding.  In _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2022, Dublin, Ireland, May 22-27, 2022_ , pages 8707–8717. Association for Computational Linguistics.
  * Li et al. (2022b) Mengze Li, Tianbao Wang, Haoyu Zhang, Shengyu Zhang, Zhou Zhao, Wenqiao Zhang, Jiaxu Miao, Shiliang Pu, and Fei Wu. 2022b.  HERO: hierarchical spatio-temporal reasoning with contrastive action correspondence for end-to-end video object grounding.  In _MM ’22: The 30th ACM International Conference on Multimedia, Lisboa, Portugal, October 10 - 14, 2022_ , pages 3801–3810. ACM.
  * Lin et al. (2012) Wan-Chen Lin, Tsung-Ting Kuo, Tung-Jia Chang, Chueh-An Yen, Chao-Ju Chen, and Shou-de Lin. 2012.  Exploiting machine learning models for chinese legal documents labeling, case classification, and sentencing prediction.  _Int. J. Comput. Linguistics Chin. Lang. Process._ , 17(4).
  * Liu et al. (2022) Dugang Liu, Weihao Du, Lei Li, Weike Pan, and Zhong Ming. 2022.  Augmenting legal judgment prediction with contrastive case relations.  In _Proceedings of the 29th International Conference on Computational Linguistics, COLING 2022, Gyeongju, Republic of Korea, October 12-17, 2022_ , pages 2658–2667. International Committee on Computational Linguistics.
  * Liu et al. (2019) Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. 2019.  Roberta: A robustly optimized BERT pretraining approach.  _CoRR_ , abs/1907.11692.
  * Luo et al. (2017) Bingfeng Luo, Yansong Feng, Jianbo Xu, Xiang Zhang, and Dongyan Zhao. 2017.  Learning to predict charges for criminal cases with legal basis.  In _Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing, EMNLP 2017, Copenhagen, Denmark, September 9-11, 2017_ , pages 2727–2736. Association for Computational Linguistics.
  * Lyu et al. (2022) Yougang Lyu, Zihan Wang, Zhaochun Ren, Pengjie Ren, Zhumin Chen, Xiaozhong Liu, Yujun Li, Hongsong Li, and Hongye Song. 2022.  Improving legal judgment prediction through reinforced criminal element extraction.  _Inf. Process. Manag._ , 59(1):102780.
  * Ma et al. (2021) Yixiao Ma, Yunqiu Shao, Bulou Liu, Yiqun Liu, Min Zhang, and Shaoping Ma. 2021.  Retrieving legal cases from a large-scale candidate corpus.  _Proceedings of the Eighth International Competition on Legal Information Extraction/Entailment, COLIEE2021_.
  * Malik et al. (2021) Vijit Malik, Rishabh Sanjay, Shubham Kumar Nigam, Kripabandhu Ghosh, Shouvik Kumar Guha, Arnab Bhattacharya, and Ashutosh Modi. 2021.  ILDC for CJPE: indian legal documents corpus for court judgment prediction and explanation.  In _Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing, ACL/IJCNLP 2021, (Volume 1: Long Papers), Virtual Event, August 1-6, 2021_ , pages 4046–4062. Association for Computational Linguistics.
  * Monroy et al. (2009) Alfredo Monroy, Hiram Calvo, and Alexander F. Gelbukh. 2009.  NLP for shallow question answering of legal documents using graphs.  In _Computational Linguistics and Intelligent Text Processing, 10th International Conference, CICLing 2009, Mexico City, Mexico, March 1-7, 2009. Proceedings_ , volume 5449 of _Lecture Notes in Computer Science_ , pages 498–508. Springer.
  * Niklaus et al. (2021) Joel Niklaus, Ilias Chalkidis, and Matthias Stürmer. 2021.  Swiss-judgment-prediction: A multilingual legal judgment prediction benchmark.  _CoRR_ , abs/2110.00806.
  * Rigoni (2014) Adam Rigoni. 2014.  Common-law judicial reasoning and analogy.  _Legal Theory_ , 20(2):133–156.
  * Saravanan et al. (2009) M. Saravanan, Balaraman Ravindran, and S. Raman. 2009.  Improving legal information retrieval using an ontological framework.  _Artif. Intell. Law_ , 17(2):101–124.
  * Shao et al. (2023) Zhenwei Shao, Zhou Yu, Meng Wang, and Jun Yu. 2023.  Prompting large language models with answer heuristics for knowledge-based visual question answering.  _CoRR_ , abs/2303.01903.
  * Shen et al. (2022) Kai Shen, Yichong Leng, Xu Tan, Siliang Tang, Yuan Zhang, Wenjie Liu, and Edward Lin. 2022.  Mask the correct tokens: An embarrassingly simple approach for error correction.  In _Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, EMNLP 2022, Abu Dhabi, United Arab Emirates, December 7-11, 2022_ , pages 10367–10380. Association for Computational Linguistics.
  * Sulea et al. (2017) Octavia-Maria Sulea, Marcos Zampieri, Shervin Malmasi, Mihaela Vela, Liviu P. Dinu, and Josef van Genabith. 2017.  Exploring the use of text classification in the legal domain.  In _Proceedings of the Second Workshop on Automated Semantic Analysis of Information in Legal Texts co-located with the 16th International Conference on Artificial Intelligence and Law (ICAIL 2017), London, UK, June 16, 2017_ , volume 2143 of _CEUR Workshop Proceedings_. CEUR-WS.org.
  * Sun (2023) Zhongxiang Sun. 2023.  A short survey of viewing large language models in legal aspect.  _CoRR_ , abs/2303.09136.
  * Touvron et al. (2023) Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurélien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. 2023.  Llama: Open and efficient foundation language models.  _CoRR_ , abs/2302.13971.
  * Wu et al. (2023) Chenfei Wu, Shengming Yin, Weizhen Qi, Xiaodong Wang, Zecheng Tang, and Nan Duan. 2023.  Visual chatgpt: Talking, drawing and editing with visual foundation models.  _CoRR_ , abs/2303.04671.
  * Wu et al. (2022) Linjuan Wu, Shaojuan Wu, Xiaowang Zhang, Deyi Xiong, Shizhan Chen, Zhiqiang Zhuang, and Zhiyong Feng. 2022.  Learning disentangled semantic representations for zero-shot cross-lingual transfer in multilingual machine reading comprehension.  In _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2022, Dublin, Ireland, May 22-27, 2022_ , pages 991–1000. Association for Computational Linguistics.
  * Wu et al. (2020) Yiquan Wu, Kun Kuang, Yating Zhang, Xiaozhong Liu, Changlong Sun, Jun Xiao, Yueting Zhuang, Luo Si, and Fei Wu. 2020.  De-biased court’s view generation with causality.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, EMNLP 2020, Online, November 16-20, 2020_ , pages 763–780. Association for Computational Linguistics.
  * Xu et al. (2020) Nuo Xu, Pinghui Wang, Long Chen, Li Pan, Xiaoyan Wang, and Junzhou Zhao. 2020.  Distinguish confusing law articles for legal judgment prediction.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, ACL 2020, Online, July 5-10, 2020_ , pages 3086–3095. Association for Computational Linguistics.
  * Yue et al. (2021) Linan Yue, Qi Liu, Binbin Jin, Han Wu, Kai Zhang, Yanqing An, Mingyue Cheng, Biao Yin, and Dayong Wu. 2021.  Neurjudge: A circumstance-aware neural framework for legal judgment prediction.  In _SIGIR ’21: The 44th International ACM SIGIR Conference on Research and Development in Information Retrieval, Virtual Event, Canada, July 11-15, 2021_ , pages 973–982. ACM.
  * Zhang et al. (2023) Wenqi Zhang, Yongliang Shen, Weiming Lu, and Yueting Zhuang. 2023.  Data-copilot: Bridging billions of data and humans with autonomous workflow.  _CoRR_ , abs/2306.07209.
  * Zhang et al. (2022) Wenqi Zhang, Kai Zhao, Peng Li, Xiao Zhu, Yongliang Shen, Yanna Ma, Yingfeng Chen, and Weiming Lu. 2022.  A closed-loop perception, decision-making and reasoning mechanism for human-like navigation.  In _Proceedings of the Thirty-First International Joint Conference on Artificial Intelligence, IJCAI-22_ , pages 4717–4724. International Joint Conferences on Artificial Intelligence Organization.  Main Track.
  * Zhao et al. (2023) Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yingqian Min, Beichen Zhang, Junjie Zhang, Zican Dong, Yifan Du, Chen Yang, Yushuo Chen, Zhipeng Chen, Jinhao Jiang, Ruiyang Ren, Yifan Li, Xinyu Tang, Zikang Liu, Peiyu Liu, Jian-Yun Nie, and Ji-Rong Wen. 2023.  A survey of large language models.  _CoRR_ , abs/2303.18223.
  * Zhong et al. (2018) Haoxi Zhong, Zhipeng Guo, Cunchao Tu, Chaojun Xiao, Zhiyuan Liu, and Maosong Sun. 2018.  Legal judgment prediction via topological learning.  In _Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, Brussels, Belgium, October 31 - November 4, 2018_ , pages 3540–3549. Association for Computational Linguistics.
  * Zhong et al. (2020) Haoxi Zhong, Chaojun Xiao, Cunchao Tu, Tianyang Zhang, Zhiyuan Liu, and Maosong Sun. 2020.  How does NLP benefit legal system: A summary of legal artificial intelligence.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, ACL 2020, Online, July 5-10, 2020_ , pages 5218–5230. Association for Computational Linguistics.

##  Appendix A Appendices

###  A.1 Top-k Accuracy

Figure 5: The top-k accuracy of CNN on CAIL dataset. Figure 6: The top-k
accuracy of CNN on CJO22 dataset. Figure 7: The top-k accuracy of BERT on CAIL
dataset. Figure 8: The top-k accuracy of BERT on CJO22 dataset.

###  A.2 More Show Cases

Figure 9: More case 1. Figure 10: More case 2. Figure 11: Case with
explanations 1. Figure 12: Case with explanations 2. Figure 13: Case with
explanations 3.


# Learning Graph Structures and Uncertainty for Accurate and Calibrated Time-
series Forecasting

Harshavardhan Kamarthi  College of Computing, Georgia Institute of Technology
USA<EMAIL_ADDRESS>,  Lingkai Kong  College of Computing, Georgia
Institute of Technology USA<EMAIL_ADDRESS>,  Alexander Rodríguez  College
of Computing, Georgia Institute of Technology USA<EMAIL_ADDRESS>,
Chao Zhang  College of Computing, Georgia Institute of Technology USA
<EMAIL_ADDRESS>and  B. Aditya Prakash  College of Computing, Georgia
Institute of Technology USA<EMAIL_ADDRESS>
###### Abstract.

Multi-variate time series forecasting is an important problem with a wide
range of applications. Recent works model the relations between time-series as
graphs and have shown that propagating information over the relation graph can
improve time series forecasting. However, in many cases, relational
information is not available or is noisy and reliable. Moreover, most works
ignore the underlying uncertainty of time-series both for structure learning
and deriving the forecasts resulting in the structure not capturing the
uncertainty resulting in forecast distributions with poor uncertainty
estimates. We tackle this challenge and introduce StoIC, that leverages
stochastic correlations between time-series to learn underlying structure
between time-series and to provide well-calibrated and accurate forecasts.
Over a wide-range of benchmark datasets StoIC provides around 16% more
accurate and 14% better-calibrated forecasts. StoIC also shows better
adaptation to noise in data during inference and captures important and useful
relational information in various benchmarks.

##  1\. Introduction

While there has been a lot of work on modeling and forecasting univariate
time-series (Hyndman and Athanasopoulos, 2018), the problem of multivariate
time-series forecasting is more challenging. This is because modeling
individual signals independently may not be sufficient to capture the
underlying relationships between the signals which are essential for strong
predictive performance. Therefore, many multivariate models model sparse
correlations between signals based on prior knowledge of underlying structure
using Convolutional networks (Li et al., 2017) or Graph Neural networks (Kipf
and Welling, 2016; Zhao et al., 2019; Yu et al., 2017).

However, in many real-world applications, the graph structure is not available
or is unreliable. In such cases, the problem of learning underlying patterns
(Franceschi et al., 2019; Shang et al., 2021; Pamfil et al., 2020) is an
active area of research (Zügner et al., 2021) in applications such as traffic
prediction and energy forecasting. Most methods use a joint learning approach
to train the parameters of both graph inference and forecasting modules.

However, most previous works focus only on point forecasting and do not
leverage uncertainty when modeling the structure. Systematically modeling this
uncertainty into the modeling pipeline can help the model adapt to unseen
patterns such as when modeling a novel pandemic (Rodriguez et al., 2021).
Therefore, the learned structure from existing models may not be adapted to
noise in data or to distributional shifts commonly encountered in real-world
datasets.

In this paper, we tackle the challenge of leveraging structure learning to
provide accurate and calibrated probabilistic forecasts for all signals of a
multivariate time-series. We introduce a novel probabilistic neural
multivariate time-series model, StoIC (Stochastic Graph Inference for
Calibrated Forecasting), that leverages functional neural process framework
(Louizos et al., 2019) to model uncertainty in temporal patterns of individual
time-series as well as a joint structure learning module that leverages both
pair-wise similarities of time-series and their uncertainty to model the graph
distribution of the underlying structure. StoIC then leverages the
distribution of learned structure to provide accurate and calibrated forecast
distributions for all the time-series.

Our contributions can be summarized as follows: (1) Deep probabilistic
multivariate forecasting model using Joint Structure learning: We propose a
Neural Process based probabilistic deep learning model that captures complex
temporal and structural correlations and uncertainty over multivariate time-
series. (2) State-of-art accuracy and calibration in multivariate forecasting:
We evaluate StoIC against previous state-of-art models in a wide range of
benchmarks and observe 16.5% more accurate and 14.7% better calibration
performance. We also show that StoIC is significantly better adapted to
provide consistent performance with the injection of varying degrees of noise
into datasets due to modeling uncertainty. (3) Mining useful structural
patterns: We provide multiple case studies to show that StoIC identifies
useful domain-specific patterns based on the graphs learned such as modeling
relations between stocks of the same sectors, location proximity in traffic
sensors, and epidemic forecasting.

##  2\. Methodology

#### Problem Formulation

Consider a multi-variate time-series dataset $\mathcal{D}$ of $N$ time-series
$\mathcal{D}=\\{\mathbf{y}_{i}\\}_{i=1}^{N}$ over $T$ time-steps. Let
$\mathbf{y}_{i}\in\mathbb{R}^{T}$ denote time-series $i$ and $y_{i}^{(t)}$ be
the value at time $t$. Further, let $\mathbf{y}^{(t)}\in\mathbb{R}^{N}$ be the
vector of all time-series values at time $t$. Given time-series values from
till current time $t$ as $\mathbf{y}^{(1:t)}$, the goal of probabilistic
multivariate forecasting is to train a model $M$ that provides a forecast
distribution:

(1) $p_{M}\left(\mathbf{y}^{(t+1:t+\tau)}|\mathbf{y}^{(1:t)};\theta\right),$

which should be accurate, i.e, has mean close to ground truth as well as
calibrated, i.e., the confidence intervals of the forecasts precisely mimic
actual empirical probabilities (Fisch et al., 2022; Kamarthi et al., 2021).

Formally, the goal of joint-structure learning for probabilistic forecasting
is to learn a global graph $G$ from $\mathbf{y}^{(1:t)}$ and leverage it to
provide accurate and well-calibrated forecast distributions:

(2)
$p_{M}\left(\mathbf{y}^{(t+1:t+\tau)}|G,\mathbf{y}^{(1:t)};\theta\right)p_{M}\left(G|\mathbf{y}^{(1:t)};\theta\right).$

Figure 1. Overview of StoIC’s pipeline. Probabilistic Time-series encoder
encodes the input time-series into stochastic latent embeddings which are used
to learn the relational graph. The relational graph is leveraged by Recurrent
Graph Neural Encoder to model both temporal and structural patterns of
multivariate time-series represented by Graph-refined embeddings. Reference
Correlation Network extracts similar patterns from past data in form of RCN
embeddings. Finally, the decoder uses Graph-refined embeddings and RCN
embeddings to derive the parameters of the forecast distribution.

#### Overview

StoIC models stochasticity and uncertainty of time-series when generating
structural relations across time-series. It also adaptively leverages
relations and uncertainty from past data using the functional process
framework (Louizos et al., 2019; Kamarthi et al., 2021). StoIC’s generative
process can be summarized as: (1) The input time-series values are encoded
using a Probabilistic Time-series Encoder (PTE) that models a multivariate
Gaussian Distribution to model each time-series capturing both time-series
patterns and inherent uncertainty. (2) The similarity between the sampled
stochastic encoding of each time-series from PTE is used to sample a graph via
the Graph Generation Module (GGM). (3) Recurrent Graph Neural Encoder (RGNE)
contains a series of Recurrent neural layers and Graph Convolutional Networks
which derive the encoding of each time-series leveraging the learned graph.
(4) We also model the similarity of encodings of input time-series with past
data using a reference correlation network (RCN). (5) Finally, StoIC uses the
graph-refined embeddings and historical information from RCN to learn the
parameters of the output distribution.

#### Probabilistic Time-series Encoder

We first model both the information and uncertainty for each of the $N$ time-
series, by using deep sequential models to capture complex temporal patterns
of the input time-series sequence $\mathbf{y}_{i}^{(t^{\prime}:t)}$. We use a
GRU (Cho et al., 2014) followed by Self-attention (Vaswani et al., 2017) over
the hidden embeddings of GRU:

(3) $\\{\mathbf{\bar{h}_{i}}\\}_{i=1}^{N}=\\{\text{Self-
Atten}(GRU(\mathbf{y}_{i}^{t^{\prime}:t}))\\}_{i=1}^{N}.$

We then model the final latent embeddings of univariate time-series as a
multivariate Gaussian distribution similar to VAEs (Kingma et al., 2019):

(4)
$\mu_{\mathbf{h}_{i}},\log\sigma_{\mathbf{h}_{i}}=NN_{h}(\mathbf{\bar{h}}_{i}),\quad\mathbf{h}_{i}\sim\mathcal{N}(\mu_{\mathbf{h}_{i}},\sigma_{\mathbf{h}_{i}}).$

where $NN_{h}$ is a single layer of feed-forward neural network. The output
latent embeddings $\mathbf{H}=\\{\mathbf{h}_{i}\\}_{i}$ are stochastic
embeddings sampled from Gaussian parameterized by
$\\{\mu(\mathbf{h}_{i}),\sigma(\mathbf{h}_{i})\\}_{i}$. which captures
uncertainty of temporal patterns.

#### Probabilistic Graph Generation Module

Since it is computationally expensive to model all possible relations between
time-series, we aim to mine sparse stochastic relations between time-series
along with the uncertainty of the underlying relations. We generate a
stochastic relational graph (SRG) $G$ of $N$ nodes that model the similarity
across all time-series. We use a stochastic process to generate $G$ leveraging
stochastic latent embeddings $\mathbf{H}$ from PTE. We parametrize the
adjacency matrix $\mathbf{A}(\mathbf{H}),\in\\{0,1\\}^{N\times N}$ of $G$ by
modelling existence of each edge $A_{i,j}$ as a Bernoulli distribution
parameterized by $\theta_{ij}$ derived as:

(5)
$\theta_{i,j}=\text{sig}(NN_{G_{2}}(NN_{G_{1}}(\mathbf{h}_{i}))+NN_{G_{1}}(\mathbf{h}_{j}))$

where $NN_{G1}$ and $NN_{G2}$ are feed-forward networks. We sample the
adjacency matrix which captures temporal uncertainty in
$\mathbf{h}_{i},\mathbf{h}_{j}$ and relational uncertainty in $\theta_{i,j}$.

(6) $(A_{i,j}=A_{j,i})\sim\text{Bernoulli}(\theta_{i,j});\forall i\leq j.$

#### Recurrent Graph Neural Encoder

We combine the relational information from SRG with temporal information via a
combination of Graph Neural Networks and recurrent networks:

(7) $\mathbf{v_{i}^{(t)}}=\text{GRU-
Cell}(\mathbf{u}_{i}^{(t-1)},y_{i}^{(t)});\\{\mathbf{u}_{i}^{(t)}\\}_{i}=GNN(\\{\mathbf{v}_{i}^{(t)}\\},A)$

We input $h_{i}$ as the initial hidden embedding for GRU-Cell at initial time-
step $t^{\prime}$ to impart temporal information from PTE. We finally combine
the intermediate embeddings $\\{\mathbf{u}_{i}^{(t)}\\}_{t=t^{\prime}}^{t}$
using self-attention to get Graph-refined Embeddings
$\mathbf{U}=\\{\mathbf{u}_{i}\\}_{i=1}^{N}$ where:

(8) $\mathbf{u}_{i}=\text{Self-Attention}(\\{\mathbf{u}_{i}^{(t)}\\}_{t}).$

#### Reference Correlation Network

This historical similarity is useful since time-series shows similar patterns
to past data. Therefore, we model relations with past historical data of all
$N$ time-series of datasets. We encode the past information of all time-series
into reference embeddings $\mathbf{G}=\\{\mathbf{g}_{j}\\}_{j=1}^{K}$:

(9)
$\mu_{\mathbf{g}_{j}},\sigma_{\mathbf{g}_{j}}=\textsc{PTE}(\mathbf{y}_{j}^{(1:t)}),\quad\mathbf{g}_{j}\sim\mathcal{N}(\mu_{\mathbf{g}_{j}},\sigma_{\mathbf{g}_{j}}).$

Then, similar to (Kamarthi et al., 2021) we sample edges of a bipartite
Reference Correlation Network $S$ between reference embeddings $\mathbf{G}$
and Graph-refined Embeddings $\mathbf{U}$ based on their similarity as:

(10)
$\kappa(\mathbf{u}_{i},\mathbf{g}_{j})=\exp(-\gamma||\mathbf{u}_{i}-\mathbf{g}_{j}||^{2}),\quad
S_{i,j}\sim\text{Bernoulli}(\kappa(\mathbf{u}_{i},\mathbf{g}_{j})).$

where $\gamma$ is a learnable parameter. To leverage the similar reference
embeddings sampled for each time-series $i$, we aggregate the sampled
reference embeddings to form the RCN embeddings
$\mathbf{Z}=\\{\mathbf{z}_{i}\\}_{i=1}^{N}$ as:

(11)
$\mathbf{z}_{i}\sim\mathcal{N}\left(\sum_{j:S_{ij}=1}NN_{z1}(\mathbf{g}_{j}),\exp(\sum_{j:S_{ij}=1}NN_{z2}(\mathbf{g}_{j}))\right)$

where $NN_{z1}$ and $NN_{z2}$ are single fully-connected layers. Therefore,
the data from reference embeddings that show similar patterns to input time-
series are more likely to be sampled.

#### Adaptive Distribution Decoder

The decoder parameterizes the forecast distribution using multiple
perspectives of information and uncertainty from previous modules: Graph-
refined embeddings $\mathbf{U}$, RCN embeddings $\mathbf{Z}$ and a global
embedding of all reference embeddings $g$ derived as: $\mathbf{g}=\text{Self-
Attention}(\\{\mathbf{g}_{j}\\}_{j}).$ However, information from each of the
modules may have varied importance based. Therefore, we use a weighted
aggregation of these embeddings:
$\mathbf{k}_{i}=l_{1}\mathbf{u}_{i}+l_{2}\mathbf{z}_{i}+l_{3}\mathbf{g}$ to
get the input embedding $\mathbf{k}_{i}$ for the decoder where
$\\{l_{i}\\}_{i=1}^{3}$ are learnable parameters. The final output forecast
distribution is derived as:

(12)
$y_{i}^{(t+1)}\sim\mathcal{N}(NN_{y1}(\mathbf{k}_{i}),\exp(NN_{y2}(\mathbf{k}_{i})))$

#### Training and Inference

The full generative pipeline of StoIC is:

$\begin{split}&P(\mathbf{y^{(t+1)}}|\mathbf{y}^{t^{\prime}:t},\mathcal{D})=\int\underbrace{P(\mathbf{H}|\mathbf{y}^{t^{\prime}:t})}_{\text{Time-
series Encoder}(\textsc{PTE})}\\\ &\underbrace{P(A|\mathbf{H})}_{\text{Graph
Generation}(\textsc{GGM})}\underbrace{P(\mathbf{U}|\mathbf{H},A)}_{\text{Recurrent
Graph Neural Encoder}(\textsc{RGNE})}\\\
&\underbrace{P(S,\\{\mathbf{g}_{j}\\}_{j}|\mathbf{U},\mathcal{D})P(\mathbf{Z}|S,\\{\mathbf{g}_{j}\\}_{j})}_{\text{Reference
Correlation
Network}(\textsc{RCN})}\underbrace{P(\mathbf{y^{(t+1)}}|\mathbf{Z,U,g})}_{\text{Decoder}}d\mathbf{H}dAd\mathbf{U}d\mathbf{g}dSd\mathbf{Z}.\end{split}$

We train the parameters of StoIC to increase the log-likelihood loss $\log
P(\mathbf{y^{(t+1)}}|\mathbf{y}^{t^{\prime}:t},\mathcal{D})$. Since
integration over high-dimensional latent random variables is intractable, we
use amortized variational inference like (Louizos et al., 2019) and construct
the variation distribution
$q(\mathbf{H,U,Z},S,A|\mathbf{y}^{t^{\prime}:t},\mathcal{D})=P(\mathbf{H}|\mathbf{y}^{t^{\prime}:t})P(A|\mathbf{H})P(S|\mathbf{U},\mathcal{D})q_{1}(\mathbf{U},\mathbf{Z}|\mathbf{y}^{t^{\prime}:t})$
where $q_{1}$ is a fully connected network over $\\{\bar{\mathbf{h}}_{i}\\}$
that parameterizes the variational distributions for $\mathbf{U}$ and
$\mathbf{Z}$. The loss is optimized using stochastic gradient descent based
Adam optimizer (Kingma and Ba, 2014). During inference, we generate Monte
Carlo samples from the full distribution
$P(\mathbf{y^{(t+1)}}|\mathbf{y}^{t^{\prime}:t},\mathcal{D})$ with discrete
sampling.

##  3\. Experiment Setup

#### Baselines

We compare with general state-of-art forecasting models that include (1)
statistical models like ARIMA (Hyndman and Athanasopoulos, 2018), (2) general
forecasting models: DeepAR (Salinas et al., 2020), EpiFNP (Kamarthi et al.,
2021) (3) Graph-learning based forecasting models: MTGNN (Wu et al., 2020),
GDN (Deng and Hooi, 2021), GTS (Shang et al., 2021), NRI (Kipf et al., 2018).
Note that we are performing probabilistic forecasting while most of the
baselines are modeled for point forecasts. Therefore, for methods like ARIMA,
DeepAR, MTGNN and GDN, we leverage an ensemble of ten independently trained
models to derive the forecast distribution (Lakshminarayanan et al., 2017).

#### Datasets

We evaluate our models against eight multivariate time-series datasets from a
wide range of applications that have been used in past works. The main
statistics of the datasets are summarized in Table 1. (1) Flu-Symptoms: We
forecast a symptomatic measure of flu incidence rate based on wILI (weighted
Influenza-like illness outpatients) that are provided by CDC for each of the
50 US states. We train on seasons from 2010 to 2015 and evaluate on seasons
from 2015 to 2020. (2) Covid-19: We forecast the weekly incidence of Covid-19
mortality from June 2020 to March 2021 for each of the 50 US states (Rodriguez
et al., 2021; Kamarthi et al., 2022) using incidence data starting from April
2020. (3) Similar to (Pamfil et al., 2020), we use the daily closing prices
for stocks of companies in S& P 100 using the yfinance package (yfi, [n. d.])
from July 2014 to June 2019. The last 400 trading days are used for testing.
(4) Electricity: We use a popular multivariate time-series dataset for power
consumption forecasting used in past works (Zügner et al., 2021). We forecast
power consumption for 15-60 minutes. We train for 1 year and test on data from
the next year. (5) Traffic prediction: We use 2 datasets related to traffic
speed prediction. METR-LA and PEMS-BAYS (Li et al., 2018) are datasets of
traffic speed at various spots in Los Angeles and San Francisco. We use the
last 10% of the time-series for testing. (6) Transit demand: NYC-Bike and NYC-
Taxi (Ye et al., 2021) measure bike sharing and taxi demand respectively in
New York from April to June 2016.

Table 1. Statistics of benchmark datasets

Dataset | Nodes | # time-steps | Sample rate | # Forecast horizon
---|---|---|---|---
Flu-Symptoms | 50 | 544 | 1 week | 4 weeks
Covid-19 | 50 | 44 | 1 week | 4 weeks
S&P-100 | 100 | 1257 | 1 day | 5 days
Electricity | 321 | 140256 | 15 min. | 60 min.
METR-LA | 207 | 34272 | 5 min. | 60 min.
PEMS-BAYS | 325 | 52116 | 5 min. | 60 min.
NYC-Bike | 250 | 4368 | 30 min. | 120 min.
NYC-Taxi | 266 | 4368 | 30 min. | 120 min.

#### Evaluation Metrics

We evaluate our model and baselines using carefully chosen metrics that are
widely used in forecasting literature.111Code:
https://anonymous.4open.science/r/Stoic_KDD24-D5A8. Supplementary:
https://anonymous.4open.science/r/Stoic_KDD24-D5A8/Struct_Supp.pdf They
evaluate for both accuracy of the mean of forecast distributions as well as
calibration of the distributions (Hyndman et al., 2006; Tsyplakov, 2013;
Jordan et al., 2017). We use RMSE for point-prediction and CRPS as well as
Confidence score (Kamarthi et al., 2021; Kuleshov et al., 2018; Fisch et al.,
2022) for measuring forecast calibrations.

#### Forecast Accuracy and Calibration

Table 2. Average Evaluation Scores (over 20 runs) for StoIC and baselines. The
best-performing scores are in bold and the second-best is underlined. StoIC
provides around 9%, 14.7%, and 16% better performance in terms of RMSE, CS and
CRPS respectively.

| Flu-Symptoms | Covid-19 | S&P-100 | Electricity
---|---|---|---|---
| RMSE | CRPS | CS | RMSE | CRPS | CS | RMSE | CRPS | CS | RMSE | CRPS | CS
ARIMA | 0.85 | 0.91 | 0.22 | 91.5 | 87.4 | 0.21 | 0.42 | 0.40 | 0.27 | 395.2 | 417.2 | 0.37
DeepAR | 0.68 | 0.97 | 0.15 | 49.1 | 48.5 | 0.19 | 0.35 | 0.41 | 0.18 | 249.9 | 267.8 | 0.18
EpiFNP | 0.64 | 0.52 | 0.05 | 43.6 | 44.5 | 0.15 | 0.22 | 0.28 | 0.16 | 257.3 | 277.0 | 0.15
MTGNN | 0.68 | 0.75 | 0.18 | 48.3 | 49.5 | 0.22 | 0.17 | 0.26 | 0.19 | 189.3 | 207.1 | 00.13
GDN | 0.73 | 0.71 | 0.15 | 42.74 | 43.2 | 0.25 | 0.18 | 0.24 | 0.24 | 249.1 | 263.6 | 0.21
NRI | 0.71 | 0.74 | 0.22 | 63.44 | 72.19 | 0.35 | 0.14 | 0.23 | 0.19 | 315.9 | 338.6 | 0.18
GTS | 0.63 | 0.66 | 0.17 | 41.96 | 36.25 | 0.22 | 0.11 | 0.18 | 0.22 | 197.4 | 208.0 | 0.16
LDS | 0.69 | 0.64 | 0.16 | 45.17 | 41.32 | 0.23 | 0.15 | 0.19 | 0.16 | 213.6 | 229.0 | 0.19
StoIC | 0.57 | 0.42 | 0.07 | 32.7 | 31.7 | 0.2 | 0.11 | 0.16 | 0.12 | 191.6 | 207.3 | 0.09
| METR-LA | PEMS-BAYS | NYC-Bike | NYC-Taxi
| RMSE | CRPS | CS | RMSE | CRPS | CS | RMSE | CRPS | CS | RMSE | CRPS | CS
ARIMA | 2.17 | 2.06 | 0.23 | 4.1 | 3.9 | 0.14 | 4.27 | 4.11 | 0.34 | 18.55 | 17.49 | 0.15
DeepAR | 2.05 | 1.93 | 0.17 | 3.8 | 4.2 | 0.16 | 3.11 | 3.15 | 0.28 | 14.52 | 15.77 | 0.17
EpiFNP | 2.11 | 1.89 | 0.14 | 4.1 | 3.9 | 0.13 | 2.94 | 2.85 | 0.23 | 15.13 | 14.97 | 0.15
MTGNN | 1.86 | 1.64 | 0.14 | 3.2 | 2.9 | 0.11 | 2.62 | 2.61 | 0.17 | 10.37 | 10.13 | 0.13
GDN | 1.95 | 1.83 | 0.16 | 3.4 | 3.1 | 0.16 | 2.74 | 1.78 | 0.15 | 11.15 | 10.64 | 0.17
NRI | 2.17 | 1.95 | 0.19 | 7.9 | 5.3 | 0.14 | 3.18 | 4.7 | 0.36 | 14.38 | 17.52 | 0.31
GTS | 1.88 | 1.74 | 0.15 | 3.7 | 3.2 | 0.13 | 2.76 | 2.85 | 0.24 | 10.75 | 12.68 | 0.18
LDS | 2.12 | 2.05 | 0.13 | 3.5 | 3.7 | 0.18 | 2.55 | 2.67 | 0.21 | 11.44 | 13.98 | 0.19
StoIC | 1.84 | 1.48 | 0.11 | 2.8 | 2.5 | 0.08 | 2.52 | 2.41 | 0.15 | 8.43 | 8.61 | 0.09

We evaluate the average performance of StoIC and all the baselines across 20
independent runs in Table 2. StoIC provides 8.8% more accurate forecasts (RMSE
scores) across all benchmarks with an impressive 16.5% more accurate forecasts
in Epidemic forecasting (Flu-Symptoms and Covid-19) and is 9% better in
traffic benchmarks. For calibration, we observe significantly better CS scores
across all benchmarks. StoIC also provides 16% higher CRPS scores over the
best-performing baseline in each task. StoIC has 34.5% and 11.5% better
performance in epidemic forecasting and traffic forecasting respectively. In
particular, we also observe that baselines like ARIMA, DeepAR, and EpiFNP
which do not learn a graph have 7-15% poorer performance in traffic benchmarks
and over 180% poorer performance in S&P-100 benchmark compared to other
baselines that learn a graph.

#### Robustness to Noise

We evaluate the efficacy of modeling the underlying structure of time-series
and learning uncertainty of each time-series in helping models adapt to noise
in the datasets. Learning a single global structure to model relations across
time-series as well as modeling uncertainty in data can help the model to
adapt to noise in datasets during inference. We therefore expect StoIC to be
further resilient to noise in data during testing. We use the models trained
on clean training time-series datasets and inject noise to time-series input
during testing. Each input time-series is first independently normalized with
0 mean and unit standard deviation and then add a gaussian noise with standard
deviation $\rho$.

Figure 2. % increase in CRPS with an increase in gaussian Noise ($\rho$) added
to the input sequence. Performance decrease of StoIC is significantly lower
than most of the baselines.

We plot the decrease in performance (measured using CRPS) with an increase in
noise, measured by $\rho$ in Figure 2. First, we observe that as we increase
$\rho$ the performance decreases for all models. We observe that the baselines
that do not learn a graph structure on average show 45-60% larger decrease in
performance compared to the rest of the baselines showing the efficacy of
structure learning for robust forecasting. However, StoIC’s decrease in
performance is significantly less compared to most of the baselines. At
$\rho=0.2$ the average performance decrease in all baselines is at 19% - 27%
whereas StoIC’s performance decrease is at 9-20%. Therefore, StoIC’s ability
to model uncertainty of time-series as well as effectively capture structural
patterns enables it to provide forecasts that are robust to noise.

#### Ablation Studies

We evaluate the efficacy of various modeling choices of StoIC. Specifically,
we access the influence of graph generation, RCN, and weighted aggregation
(Equation 2) via ablation studies. StoIC on average outperforms all the
ablation variants with 8.5% better accuracy and 3.2% better CS. Graph
generation is the most impactful for performance followed by RCN. For
additional details see Appendix C.

#### Relations Captured by Inferred Graphs

We consider various meaningful domain-specific relations for the time-series
of these datasets and study how well StoIC’s inferred graphs capture them. We
observed that StoIC inferred strong relations between time-series of stocks of
the same sectors in S&P-100. In the case of PEMS-BAYS and METR-LA, the graph
inferred by StoIC between traffic sensors is most highly correlated to actual
proximity of sensors with each other compared to other baselines. Finally, the
graphs inferred in Flu-Symptoms and Covid-19 are correlated with geographical
adjacency of regions and road density of connecting regions. We provide more
details in Appendix Section C.1

##  4\. Conclusion

We introduced StoIC, a probabilistic multivariate forecasting model that
performs joint structure learning and forecasting leveraging uncertainty in
time-series data to provide accurate and well-calibrated forecasts. We
observed that StoIC performs 8.5% better in accuracy and 14.7% better in
calibration over multivariate time-series benchmarks from various domains. Due
to structure learning and uncertainty modeling, we also observed that StoIC
better adapts to the injection of varying degrees of noise to data during
inference with StoIC’s performance drop being almost half of the other state-
of-art baselines. Finally, we observed that StoIC identifies useful patterns
across time-series based on inferred graphs such as the correlation of stock
from similar sectors, location proximity of traffic sensors and geographical
adjacency and mobility patterns across US states for epidemic forecasting.
While our work focuses on time-series with real values modeled as Gaussian
distribution, our method can be extended to other distributions modeling
different kinds of continuous signals. Further, StoIC only models a single
global structure across time-series similar to most previous works. Therefore,
extending our work to learn dynamic graphs that can adapt to changing
underlying time-series relations or model multiple temporal scales could be an
important direction for future research.

## References

  * (1)
  * yfi ([n. d.]) [n. d.].  yfinance · PyPI.  https://pypi.org/project/yfinance/.  (Accessed on 10/12/2022).
  * Cho et al. (2014) Kyunghyun Cho, Bart Van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. 2014.  On the properties of neural machine translation: Encoder-decoder approaches.  _arXiv preprint arXiv:1409.1259_ (2014).
  * Deng and Hooi (2021) Ailin Deng and Bryan Hooi. 2021.  Graph neural network-based anomaly detection in multivariate time series. In _Proceedings of the AAAI Conference on Artificial Intelligence_ , Vol. 35. 4027–4035.
  * Fisch et al. (2022) Adam Fisch, Tommi Jaakkola, and Regina Barzilay. 2022.  Calibrated Selective Classification.  _arXiv preprint arXiv:2208.12084_ (2022).
  * Franceschi et al. (2019) Luca Franceschi, Mathias Niepert, Massimiliano Pontil, and Xiao He. 2019.  Learning discrete structures for graph neural networks. In _International conference on machine learning_. PMLR, 1972–1982.
  * Hyndman et al. (2006) Rob J Hyndman et al. 2006\.  Another look at forecast-accuracy metrics for intermittent demand.  _Foresight: The International Journal of Applied Forecasting_ 4, 4 (2006), 43–46.
  * Hyndman and Athanasopoulos (2018) Rob J Hyndman and George Athanasopoulos. 2018.  _Forecasting: principles and practice_.  OTexts.
  * Jang et al. (2016) Eric Jang, Shixiang Gu, and Ben Poole. 2016.  Categorical reparameterization with gumbel-softmax.  _arXiv preprint arXiv:1611.01144_ (2016).
  * Jordan et al. (2017) Alexander Jordan, Fabian Krüger, and Sebastian Lerch. 2017.  Evaluating probabilistic forecasts with scoringRules.  _arXiv preprint arXiv:1709.04743_ (2017).
  * Kamarthi et al. (2021) Harshavardhan Kamarthi, Lingkai Kong, Alexander Rodríguez, Chao Zhang, and B Aditya Prakash. 2021.  When in doubt: Neural non-parametric uncertainty quantification for epidemic forecasting.  _Advances in Neural Information Processing Systems_ 34 (2021), 19796–19807.
  * Kamarthi et al. (2022) Harshavardhan Kamarthi, Lingkai Kong, Alexander Rodríguez, Chao Zhang, and B Aditya Prakash. 2022.  CAMul: Calibrated and Accurate Multi-view Time-Series Forecasting. In _Proceedings of the ACM Web Conference 2022_. 3174–3185.
  * Kingma and Ba (2014) Diederik P Kingma and Jimmy Ba. 2014.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ (2014).
  * Kingma et al. (2019) Diederik P Kingma, Max Welling, et al. 2019\.  An introduction to variational autoencoders.  _Foundations and Trends® in Machine Learning_ 12, 4 (2019), 307–392.
  * Kipf et al. (2018) Thomas Kipf, Ethan Fetaya, Kuan-Chieh Wang, Max Welling, and Richard Zemel. 2018.  Neural relational inference for interacting systems. In _International Conference on Machine Learning_. PMLR, 2688–2697.
  * Kipf and Welling (2016) Thomas N Kipf and Max Welling. 2016.  Semi-supervised classification with graph convolutional networks.  _arXiv preprint arXiv:1609.02907_ (2016).
  * Krishnan et al. (2017) Rahul Krishnan, Uri Shalit, and David Sontag. 2017.  Structured inference networks for nonlinear state space models. In _Proceedings of the AAAI Conference on Artificial Intelligence_ , Vol. 31.
  * Kuleshov et al. (2018) Volodymyr Kuleshov, Nathan Fenner, and Stefano Ermon. 2018.  Accurate uncertainties for deep learning using calibrated regression. In _International cuonference on machine learning_. PMLR, 2796–2804.
  * Lakshminarayanan et al. (2017) Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. 2017.  Simple and scalable predictive uncertainty estimation using deep ensembles.  _Advances in neural information processing systems_ 30 (2017).
  * Li et al. (2021) Longyuan Li, Junchi Yan, Xiaokang Yang, and Yaohui Jin. 2021.  Learning interpretable deep state space model for probabilistic time series forecasting.  _arXiv preprint arXiv:2102.00397_ (2021).
  * Li et al. (2017) Yaguang Li, Rose Yu, Cyrus Shahabi, and Yan Liu. 2017.  Diffusion convolutional recurrent neural network: Data-driven traffic forecasting.  _arXiv preprint arXiv:1707.01926_ (2017).
  * Li et al. (2018) Yaguang Li, Rose Yu, Cyrus Shahabi, and Yan Liu. 2018.  Diffusion Convolutional Recurrent Neural Network: Data-Driven Traffic Forecasting. In _International Conference on Learning Representations_.
  * Louizos et al. (2019) Christos Louizos, Xiahan Shi, Klamer Schutte, and Max Welling. 2019.  The functional neural process.  _Advances in Neural Information Processing Systems_ 32 (2019).
  * Nelson and Rae (2016) Garrett Dash Nelson and Alasdair Rae. 2016.  An economic geography of the United States: From commutes to megaregions.  _PloS one_ 11, 11 (2016), e0166083.
  * Pamfil et al. (2020) Roxana Pamfil, Nisara Sriwattanaworachai, Shaan Desai, Philip Pilgerstorfer, Konstantinos Georgatzis, Paul Beaumont, and Bryon Aragam. 2020.  Dynotears: Structure learning from time-series data. In _International Conference on Artificial Intelligence and Statistics_. PMLR, 1595–1605.
  * Rangapuram et al. (2018) Syama Sundar Rangapuram, Matthias W Seeger, Jan Gasthaus, Lorenzo Stella, Yuyang Wang, and Tim Januschowski. 2018.  Deep state space models for time series forecasting.  _Advances in neural information processing systems_ 31 (2018).
  * Rodriguez et al. (2021) Alexander Rodriguez, Anika Tabassum, Jiaming Cui, Jiajia Xie, Javen Ho, Pulak Agarwal, Bijaya Adhikari, and B Aditya Prakash. 2021.  Deepcovid: An operational deep learning-driven framework for explainable real-time covid-19 forecasting. In _Proceedings of the AAAI Conference on Artificial Intelligence_ , Vol. 35. 15393–15400.
  * Rosensteel et al. (2021) Grant E Rosensteel, Elizabeth C Leee, Vittoria Colizza, and Shweta Bansal. 2021.  Characterizing an epidemiological geography of the United States: influenza as a case study.  _medRxiv_ (2021).
  * Salinas et al. (2020) David Salinas, Valentin Flunkert, Jan Gasthaus, and Tim Januschowski. 2020.  DeepAR: Probabilistic forecasting with autoregressive recurrent networks.  _International Journal of Forecasting_ 36, 3 (2020), 1181–1191.
  * Seo et al. (2018) Youngjoo Seo, Michaël Defferrard, Pierre Vandergheynst, and Xavier Bresson. 2018.  Structured sequence modeling with graph convolutional recurrent networks. In _International conference on neural information processing_. Springer, 362–373.
  * Shang et al. (2021) Chao Shang, Jie Chen, and Jinbo Bi. 2021.  Discrete graph structure learning for forecasting multiple time series.  _arXiv preprint arXiv:2101.06861_ (2021).
  * Tsyplakov (2013) Alexander Tsyplakov. 2013.  Evaluation of probabilistic forecasts: proper scoring rules and moments.  _Available at SSRN 2236605_ (2013).
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  _Advances in neural information processing systems_ 30 (2017).
  * Wu et al. (2020) Zonghan Wu, Shirui Pan, Guodong Long, Jing Jiang, Xiaojun Chang, and Chengqi Zhang. 2020.  Connecting the dots: Multivariate time series forecasting with graph neural networks. In _Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining_. 753–763.
  * Ye et al. (2021) Junchen Ye, Leilei Sun, Bowen Du, Yanjie Fu, and Hui Xiong. 2021.  Coupled layer-wise graph convolution for transportation demand prediction. In _Proceedings of the AAAI conference on artificial intelligence_ , Vol. 35. 4617–4625.
  * Yu et al. (2017) Bing Yu, Haoteng Yin, and Zhanxing Zhu. 2017.  Spatio-temporal graph convolutional networks: A deep learning framework for traffic forecasting.  _arXiv preprint arXiv:1709.04875_ (2017).
  * Zhao et al. (2019) Ling Zhao, Yujiao Song, Chao Zhang, Yu Liu, Pu Wang, Tao Lin, Min Deng, and Haifeng Li. 2019.  T-gcn: A temporal graph convolutional network for traffic prediction.  _IEEE Transactions on Intelligent Transportation Systems_ 21, 9 (2019), 3848–3858.
  * Zügner et al. (2021) Daniel Zügner, François-Xavier Aubet, Victor Garcia Satorras, Tim Januschowski, Stephan Günnemann, and Jan Gasthaus. 2021.  A study of joint graph inference and forecasting.  _arXiv preprint arXiv:2109.04979_ (2021).

Supplementary for the paper "Learning Latent Graph Structures for Accurate and
Calibrated Time-series Forecasting"

We run our models on an Nvidia Tesla V100 GPU and found that it takes less
than 4 GB of memory for all benchmarks. The code for our model is available at
an anonymized repository https://anonymous.4open.science/r/Stoic_KDD24-D5A8
and will be released publicly when accepted.

##  Appendix A Related Work

#### Multivariate forecasting using domain-based structural data

Deep neural networks have achieved great success in probabilistic time series
forecasting. DeepAR (Salinas et al., 2020) trains an auto-regressive recurrent
network to predict the parameters of the forecast distributions. Other works
including deep Markov models (Krishnan et al., 2017) and deep state space
models (Rangapuram et al., 2018; Li et al., 2021) explicitly model the
transition and emission components with neural networks. Recently, EpiFNP
(Kamarthi et al., 2021) leverages functional neural processes and achieves
state-of-art performance in epidemic forecasting. However, all these methods
treat each time series individually and suffer from limited forecasting
performance. Leveraging the relation structure among time-series to improve
forecasting performance is an emerging area. GCRN (Seo et al., 2018), DCRNN
(Li et al., 2018), STGCN (Yu et al., 2017) and T-GCN (Zhao et al., 2019) adopt
graph neural networks to capture the relationships among time series and
provide better representations for each individual sequence. However, these
methods all assume that the ground-truth graph structure is available in
advance, which is often unknown in many real world applications.

#### Structure learning for time-series forecasting

When the underlying structure is unknown, we need to jointly perform graph
structure learning and time-series forecasting. MTGNN (Wu et al., 2020) and
GDN (Deng and Hooi, 2021) parameterize the graph as a degree-k graph to
promote sparsity but their training can be difficult since the top-K operation
is not differentiable. GTS (Shang et al., 2021) uses the Gumbel-softmax trick
(Jang et al., 2016) for differentiable structure learning and uses prior
knowledge to regularize the graph structure. The graph learned by GTS is a
global graph shared by all the time series. Therefore, it is not flexible
since it cannot adjust the graph structure for changing inputs at inference
time. NRI (Kipf et al., 2018) employs a variational auto-encoder architecture
and can produce different structures for different encoding inputs. It is more
flexible than GTS but needs more memory to store the individual graphs.
However, as previously discussed, these works do not model the uncertainty of
time-series during structure learning and forecasting and do not focus on the
calibration of their forecast distribution.

##  Appendix B Training details

The architecture of PTE and RCN is similar to (Kamarthi et al., 2021) with GRU
being bi-directional and having 60 hidden units.
$NN_{h},NN_{G1},NN_{G2},NN_{z1}$ and $NN_{z2}$ and GRU of RGNE also have 60
hidden units. Therefore, Graph-refined embeddings $\mathbf{U}$, RCN embeddings
$\mathbf{Z}$ and global embedding $\mathbf{g}$ are 60 dimensional vectors.

We used Adam optimizer (Kingma and Ba, 2014) with a learning rate of 0.001. We
found that using a batch size of 64 or 128 provided good performance with
stable training. We used early stopping with patience of 200 epochs to prevent
overfitting. For each of the 20 independent runs, we initialized the random
seeds of all packages to 1-20. In general, variance in performance across
different random seeds was not significant for all models.

##  Appendix C Ablation Studies

We evaluate the efficacy of various modeling choices of StoIC. Specifically,
we access the influence of graph generation, RCN, and weighted aggregation
(Equation 2) via the following ablation variants of StoIC:

  * •

StoIC-NoGraph: We remove the GGM and GCN modules of RGNE and therefore do not
use Graph-refined embeddings $\mathbf{U}$ in the decoder.

  * •

StoIC-NoRefCorr: We remove the RCN module and do not use RCN embeddings
$\mathbf{Z}$ in the decoder.

  * •

StoIC-NoWtAgg: We replace weighted aggregation with concatenation in Equation
2 as

(13) $\mathbf{k}_{i}=\mathbf{u}_{i}\oplus\mathbf{z}_{i}\oplus\mathbf{g}$

where $\oplus$ is the concatenation operator.

Table 3. Average Evaluation Scores (over 20 runs) for StoIC and ablation
variants.

| Flu-Symptoms | Covid-19 | S&P-100 | Electricity
---|---|---|---|---
| RMSE | CRPS | CS | RMSE | CRPS | CS | RMSE | CRPS | CS | RMSE | CRPS | CS
StoIC | 0.57 | 0.42 | 0.07 | 32.7 | 31.7 | 0.2 | 0.11 | 0.16 | 0.12 | 191.6 | 207.3 | 0.09
StoIC-NoGraph | 0.67 | 0.51 | 0.06 | 42.7 | 41.6 | 0.11 | 0.24 | 0.31 | 0.16 | 226.5 | 248.5 | 0.13
StoIC-NoRefCorr | 0.65 | 0.68 | 0.15 | 44.72 | 40.66 | 0.19 | 0.12 | 0.21 | 0.15 | 196.4 | 211.0 | 0.11
StoIC-NoWtAgg | 0.63 | 0.44 | 0.06 | 31.6 | 28.3 | 0.19 | 0.13 | 0.19 | 0.18 | 217.5 | 237.6 | 0.18
| METR-LA | PEMS-BAYS | NYC-Bike | NYC-Taxi
| RMSE | CRPS | CS | RMSE | CRPS | CS | RMSE | CRPS | CS | RMSE | CRPS | CS
StoIC | 1.84 | 1.48 | 0.11 | 2.8 | 2.5 | 0.08 | 2.52 | 2.41 | 0.15 | 8.43 | 8.61 | 0.09
StoIC-NoGraph | 2.19 | 1.85 | 0.17 | 3.9 | 4 | 0.15 | 3.41 | 3.16 | 0.19 | 13.65 | 12.89 | 0.21
StoIC-NoRefCorr | 2.25 | 2.18 | 0.15 | 3.4 | 3.1 | 0.12 | 2.86 | 2.72 | 0.17 | 10.53 | 9.24 | 0.14
StoIC-NoWtAgg | 1.86 | 1.54 | 0.14 | 3.7 | 3.6 | 0.09 | 2.95 | 2.69 | 0.21 | 9.25 | 8.37 | 0.11

#### Forecasting performance

As shown in Table 3, StoIC on average outperforms all the ablation variants
with 8.5% better accuracy and 3.2% better CS. On average, the worst performing
variant is StoIC-NoGraph followed by StoIC-NoRefCorr and StoIC-NoWtAgg,
showing the importance of graph generation and leveraging learned relations
across time-series.

(a) Flu-Symptoms

(b) S&P-100

(c) METR-LA

(d) PEMS-BAYS

(e) Covid-19

(f) Electricity

Figure 3. StoIC is better adapted to noise than other ablation variants. %
increase in CRPS of StoIC and ablation variants with an increase in gaussian
Noise ($\rho$) added to the input sequence.

#### Robustness to noise

Similar to Section 3, we also test the ablation variants’ robustness in the
NFI task with performance decrease compared in Figure 3. Performance
continuously decreases with an increase in $\rho$. StoIC-NoGraph is most
susceptible to a decrease in performance with a 22-41% decrease in performance
at $\rho=0.2$ which again shows the importance of structure to robustness.
StoIC-NoRefCorr and StoIC-NoWtAgg’s performance degradation range from 15-32%.
Finally, we observe that StoIC is again significantly more resilient to noise
compared to all other variants.

###  C.1. Relations Captured by Inferred Graphs

As mentioned before, for all benchmarks there is no ‘ground-truth’ structure
to compare against. However, in line with previous works (Shang et al., 2021;
Pamfil et al., 2020; Wu et al., 2020; Rosensteel et al., 2021), we consider
various meaningful domain-specific relations for the time-series of these
datasets and study how well StoIC’s inferred graphs capture them through
following case studies. Note that, for StoIC and baselines that use a sampling
strategy to construct the graph (GTS, NRI, GTS, LDS), we calculate edge
probability of each pair of nodes based on sampled graphs. For other graph
generating baselines, we directly use the graphs inferred.

####  Case Study 1: Sector-level correlations of stocks in S&P-100

Two stocks representing companies from the same sectors typically influence
each other more than stocks from different sectors. Therefore, we measure the
correlation of edge probabilities for stocks in the same sectors. We first
construct a Sector-partion graph as follows. We partition the stock time-
series into sectors and construct a set of fully connected components where
each component contains nodes of a sector. There are no edges across different
sectors. We then measure the correlation of the edge probability matrix with
the adjacency matrix of Sector-partion graph.

We observed a strong correlation score of 0.73 for graphs generated by StoIC.
This was followed by graphs from GTS and LDS with correlation scores of 0.67
and 0.55. Other baselines’ graphs provided poor correlation scores below 0.35.
Interestingly, this trend also correlated with the performance in forecasting
with StoIC, GTS and LDS being the top three best-performing models. We also
observed that the correlation score of StoIC-NoRefCorr was similar to StoIC
whereas StoIC-NoWtAgg also provided low correlation scores (0.39) though their
forecasting performances are comparable.

Table 4. Average edge correlation of generated graphs from the models with Sector-partition graph (for S&P-100) and proximity graph (for METR-LA and PEMS-BAYS). | MTGNN | GDN | NIR | GTS | LDS | StoIC
---|---|---|---|---|---|---
S&P-100 | 0.05 | 0.42 | 0.34 | 0.67 | 0.55 | 0.73
METR-LA | 0.01 | 0.23 | 0.16 | 0.27 | 0.19 | 0.23
PEMS-BAYS | 0 | 0.28 | 0.21 | 0.7 | 0.25 | 0.61

####  Case Study 2: Identifying location proximity in traffic forecasting
benchmarks

Since sensors that are located close to together may have a larger probability
of showing similar or correlated signals, we study if the generated graphs
capture information about the proximity of the location of the traffic
sensors. We use the road location information of time-series in METR-LA and
PEMS-BAYS datasets and construct the proximity graph based on pairwise road
distances similar to (Zügner et al., 2021). Note that we did not feed any
location-based information to the models during training.

We again measure the similarity of generated graphs with the proximity graph.
We observe that GTS and StoIC provide the strongest correlation scores for
both PEMS-BAYS and METR-LA datasets with average scores of 0.7 and 0.61 for
PEMS-BAYS and 0.27 and 0.23 for METR-LA. Due to the lower correlation scores
for METR-LA, the proximity of sensors may not be useful in modeling relations
across time-series. Comparing with ablation variants, we observed that both
StoIC-NoRefCorr and StoIC-NoWtAgg showed similar correlation scores to StoIC.

####  Case Study 3: Inferring geographical adjacency and mobility for Epidemic
Forecasting

(a) Covid-19

(b) Flu-Symptoms

Figure 4. Most confident edges ($\theta>0.8$) of state relation graph
generated by StoIC for Epidemic Forecasting tasks. Note that most edges
connect US states which have geographical proximity.

We observe the most confident edges generated by StoIC for Covid-19 and Flu-
Symptoms tasks and find that most of the edges map to adjacent states or
states with strong geographical proximity similar to past works (Rosensteel et
al., 2021). Further, we observed that the specific states on which the graph
relations are most confident are also connected by a higher density of roads
with frequent commutes (Nelson and Rae, 2016). This shows that StoIC can go
beyond simple patterns and infer complex mobility patterns across states
leveraging past epidemic incidence data. Hence, StoIC exploits the useful
relations pertaining to both geographical adjacency and mobility across these
states to provide state-of-art forecasting performance.