From Local Recovery to Unified Generation: Discovering Triple Set in Knowledge Graphs via Observation-Anchored Latent Diffusion

CCS Concepts

• Computing methodologies → Knowledge representation and reasoning; Semantic networks; Reasoning about belief and knowledge.

Introduction

Knowledge Graphs (KGs) have become fundamental infrastructures for a wide range of intelligent applications [14, 35, 39], including recommendation systems [20], question answering [16, 25], information retrieval [18, 23], and scientific discovery [8, 11, 46]. Despite their broad applicability, real-world KGs are inherently incomplete, motivating extensive research on knowledge graph completion (KGC) [10, 43]. Traditional KGC methods mainly focus on predicting missing links independently through embedding-based scoring functions [2, 3, 24, 45] or graph neural reasoning mechanisms [7, 51]. However, in many scenarios, the goal is not to predict a single missing triple, but to recover a set of missing triples that complete the

observed graph [50]. Thus, Triple Set Prediction (TSP) has recently emerged as a promising direction that achieves KG completion by predicting an entire set of missing triples conditioned on the observed graph [17, 49]. Compared with conventional link prediction [5, 6], TSP emphasizes structured generation and holistic graph completion. TSP can also be formulated as conditional graph generation, a perspective explored by DiffTSP [17].

Existing TSP methods predominantly follow a subgraph-based completion paradigm. As illustrated in Figure 1(a), these approaches first partition the original knowledge graph into local subgraphs and then perform completion in each subgraph [17, 50]. While such a paradigm reduces computational complexity and improves local reasoning capability, this computational simplification can limit candidate coverage and access to cross-subgraph context. First, as shown in Figure 1, it breaks cross-subgraph dependencies, so that a target relation such as the connection between nodes 2 and 3 cannot be recovered by any isolated local completion model. Second, partitioning may discard boundary context. For example, removing the edge between nodes 2 and 1 makes node 2 locally ambiguous and leads to a spurious prediction between nodes 2 and 4. More broadly, independently completed subgraphs lack a unified global latent space, which weakens graph-level semantic coordination and can yield missing or inconsistent predictions after merging. These limitations motivate a central question: can we recover missing triples across the whole knowledge graph without restricting prediction to subgraphs, while keeping generation computationally manageable? Directly generating a dense multi-relational adjacency remains costly, as its size grows quadratically with the number of entities and linearly with the number of relation types. Our key idea is to compress the state on which generation operates, rather than restrict the prediction domain through graph partitioning. To this end, we propose unified latent generation. As illustrated in Figure 1(b), this paradigm performs graph completion in a compact latent space, allowing triples to be generated without requiring their endpoints to co-occur in a local subgraph.

While unified latent modeling avoids the structural fragmentation caused by graph partitioning, recovering the target KG still requires transforming the observed representation into the target one. To instantiate this paradigm, we use conditional latent diffusion to learn recovery from perturbed target representations in TSP. Diffusion performs generation through iterative refinement [1, 30, 31], which is suitable for TSP because missing triples are structurally correlated. The denoising process progressively refines the latent graph representation and adjusts correlated relational patterns among triples [15, 38]. However, there are two challenges in applying a generic latent diffusion model to TSP. First, a straightforward VQ-VAE [40] is not sufficient for TSP since KGs are sparse and multi-relational [9, 32, 47], while standard VQ-VAE is not designed to preserve relational adjacency patterns. Second, standard diffusion [19] is not ideal for TSP since the observed KG already provides a strong structural anchor. Standard diffusion gradually perturbs the latent representation toward an isotropic Gaussian prior, which may unnecessarily destroy the structural information already contained in the observed graph.

To address these challenges, we propose LatentDiff, a unified conditional latent generative model for TSP. For the first challenge, we propose a relation-aware residual-quantized knowledge graph autoencoder. Its encoder constructs latent embeddings from relational adjacency and uses relational view attention to capture dependencies among different relational contexts. To improve the expressiveness of latent space, we use chunked residual vector quantization, which decomposes each latent vector into multiple chunks and quantizes them through multi-level residual vector quantization progressively to enhance reconstruction capacity. Finally, a factorized triple decoder injects relational reasoning bias into KG reconstruction.

For the second challenge, we develop an observation-anchored latent graph diffusion module. Building upon the mean-reverting diffusion principle, its forward process gradually perturbs the target latent representation toward the observed latent. The reverse process is parameterized by a dual-view latent graph denoiser, which captures dependencies along both the relation and entity axes and progressively reconstructs the missing relational information. Our main contributions include: (1) We formulate knowledge graph triple set prediction as a unified conditional latent generation problem, avoiding the candidate space and contextual restrictions induced by subgraph partitioning. (2) We propose a relation-aware residual-quantized knowledge graph autoencoder to construct a compact shared latent space for reconstructing relation-specific adjacency patterns. (3) We introduce an observation-anchored latent graph diffusion model that combines a mean-reverting diffusion process with a dual-view latent graph denoiser to progressively recover missing relational information in the latent space. (4) Extensive experiments demonstrate that our method outperforms state-of-the-art baselines across multiple datasets.

Related Work

Triple Set Prediction

Triple set prediction extends conventional KGC from predicting individual triples to recovering a set of missing triples from an observed KG and provides a more comprehensive setting for automatic KGC. Recent studies address this task from various angles: GPHT [50] introduces dedicated evaluation protocols and metrics for this task, LLM TSP [49] and SMuRT [48] leverage an LLM to generate the triple set, and DiffTSP [17] designs a discrete diffusion model for TSP. To make set-level prediction tractable on KGs, existing models often decompose the original graph into smaller subgraphs and perform local subgraph prediction. However, subgraph partitioning suffers from separating target entities and fragmenting the relational evidence required to predict the connecting triples.

Diffusion Models for Graphs

Diffusion models have demonstrated remarkable success in continuous domains such as image synthesis [26, 29, 36, 44] and sequential generation [34]. Several recent works [21, 33, 42, 53] attempt to extend diffusion models to graphs by performing iterative denoising over graph structures or learned representations [21, 42]. Recent studies further perform diffusion in compact latent spaces. LGD [53] develops an encoder decoder latent diffusion framework for graph generation and prediction, while HypDiff [12] incorporates geometry-aware latent diffusion to better preserve graph topology. These studies mostly focus on generic graph generation or prediction. Triple set prediction instead requires recovering missing relational facts from an observed multi-relational graph, where the latent space must preserve relation-specific structures while supporting conditional graph completion. Different from these studies, we focus on conditional triple set recovery in multi-relational KGs. Beyond graph generation, diffusion models have also been explored for knowledge graph completion. For example, some methods use diffusion to learn triple embeddings for KGC [27, 28], while others use diffusion as the hard negative samples generator [22]. However, these methods primarily evaluate link prediction, whereas TSP requires selecting a set of missing triples from the candidate space defined by the observed KG.

Problem Definition

Formally, a knowledge graph is usually represented as $\mathcal{G} = (\mathcal{E}, \mathcal{R}, \mathcal{T})$, where $\mathcal{E}$ and $\mathcal{R}$ denote the entity set and relation set, respectively, and $\mathcal{T} \subseteq\mathcal{E} \times\mathcal{R} \times\mathcal{E}$ is the set of factual triples. Each triple $(h,r,t) \in\mathcal{T}$ describes a relation $r$ from a head entity $h$ to a tail entity $t$. In TSP, we are given an observed incomplete knowledge graph $\mathcal{G}_{\mathrm{obs}} = (\mathcal{E}, \mathcal{R}, \mathcal{T}_{\mathrm{obs}})$, where $\mathcal{T}_{\mathrm{obs}}$ is a subset of the triples in the target graph. The goal is to recover a set of missing triples $\mathcal{T}_{\mathrm{miss}} = \mathcal{T}_{\mathrm{tar}} \setminus\mathcal{T}_{\mathrm{obs}}$, where $\mathcal{G}_{\mathrm{tar}} = (\mathcal{E}, \mathcal{R}, \mathcal{T}_{\mathrm{tar}})$ denotes the target KG. The predicted KG can then be written as

where $\widehat{\mathcal{T}}_{\mathrm{miss}}$ is the predicted missing triple set. The learning objective is to generate a triple set $\widehat{\mathcal{T}}_{\mathrm{miss}}$ that approximates the ground-truth missing set $\mathcal{T}_{\mathrm{miss}}$. During inference, the predicted triples that already appear in $\mathcal{T}_{\mathrm{obs}}$ are filtered out, and the final output is evaluated against the missing set.

Method

Overview

As illustrated in Figure 2, LatentDiff is a two-stage framework that first learns a compact KG latent space and then performs conditional generation within it. In Stage I, a relation-aware residual-quantized knowledge graph autoencoder (details in § 3.2), consisting of a relation-aware graph encoder, a chunked residual vector quantizer, and a factorized triple decoder, maps a KG to $Z \in\mathbb{R}^{|\mathcal{E}|\times|\mathcal{R}|\times d}$, where each embedding $z_{h,r}$ represents a head-relation context for reconstructing its tail entities. In Stage II, an observation-anchored latent graph diffusion model generates the target latent conditioned on the observed latent (details in § 3.3). The forward corruption uses the observed representation as its mean-reversion reference, and iterative latent recovery produces a continuous estimate for triple decoding.

Relation-aware Residual-Quantized KG AutoEncoder

The goal of LatentDiff’s first stage (Stage I) is to learn a compact and reconstructive latent space for KGs. To better preserve relation-specific adjacency patterns, we design a relation-aware residual-quantized KG autoencoder.

Variable-Density KG Construction

Stage I is trained as an autoencoder, where the input graph itself serves as the reconstruction target. To expose the autoencoder to KGs with different sparsity levels, we randomly sparsify the training graph at each iteration. For each training iteration, we sample a ratio $\rho\sim\mathcal{U}(0,1)$, and construct a masked KG $\mathcal{G}_{\rho}=(\mathcal{E},\mathcal{R},\mathcal{T}_{\rho})$ by independently keeping each positive triple in the training graph with probability $\rho$. The autoencoder takes $\mathcal{G}_{\rho}$ as input and is trained to reconstruct the input.

Relation-aware Graph Encoder

For each head entity $h$ and relation $r$, we define a tail indicator vector $a_{h,r}\in\{0,1\}^{|\mathcal{E}|}$, where the $t$-th entry is 1 if $(h,r,t)$ is in the KG and 0 otherwise. The relation-aware graph encoder maps each sparse tail indicator vector into a dense latent embedding: $x_{h,r}=W_a a_{h,r}$, where $W_a$ is a projection matrix. For each head entity $h$, we collect its relation-specific latent as $X_h$:

To model dependencies among different relational views, we propose relational view attention:

where attention is performed along the relation dimension. The output is combined with the original projected representation through a residual connection: $X_h=X_h+\widetilde{X}_h$. Stacking all head entities gives the encoder output $X=E_{\phi}(\mathcal{G}_{\rho})\in\mathbb{R}^{|\mathcal{E}|\times|\mathcal{R}|\times d}$, where $E_{\phi}$ is the encoder.

Chunked Residual Vector Quantization

The continuous encoder output is further discretized by a chunked residual vector quantizer. Specifically, we decompose each latent embedding into $K$ chunks: $x_{h,r}=[x_{h,r}^{(1)},x_{h,r}^{(2)},\ldots,x_{h,r}^{(K)}]$. For each chunk, we apply residual vector quantization with $L_c$ codebooks. Let $r^{(k,1)}=x_{h,r}^{(k)}$ be the initial residual of the $k$-th chunk. At the $\ell$-th residual quantization level, the nearest codeword is selected as

where $e_j^{(k,\ell)}$ is the $j$-th codeword in the $\ell$-th codebook of the $k$-th chunk. The residual is then updated by

The quantized representation of the $k$-th chunk is

Finally, the quantized latent embedding is obtained by concatenating all quantized chunks: $z_{q,h,r}=[q_{h,r}^{(1)},q_{h,r}^{(2)},\ldots,q_{h,r}^{(K)}]$. We adopt the straight-through estimator for the non-differentiable nearest-neighbor assignment:

which equals $z_{q,h,r}$ in the forward pass while propagating gradients to $x_{h,r}$ in the backward pass. $z_{h,r}$ is used as the decoder input and $\operatorname{sg}(\cdot)$ is the stop-gradient operator. The chunk-level VQ training objective is

where $\Omega$ denotes the set of head-relation pairs, $K$ is the number of chunks, and $\lambda_m$ is the commitment weight.

Factorized Triple Decoder

The decoder reconstructs triples from the quantized latent embeddings. For a triple $(h,r,t)$, the factorized scoring function is defined as

where $z_{h,r}$ is the quantized latent representation of the head-relation pair $(h,r)$, $e_t$ is the tail entity embedding, $U(\cdot)$ is the projection function, and $b_r$ is relation bias. The probability of a triple being positive is then $p(h,r,t)=\sigma(s(h,r,t))$.

Training Objective

We train the relation-aware residual-quantized KG autoencoder with negative sampling. For each positive triple $(h,r,t)\in\mathcal{T}_{\rho}$, we sample $K_{\mathrm{neg}}$ negative triples by corrupting the tail entity:

The reconstruction loss is

The overall Stage-I objective is

Here, $\mathcal{L}_{\mathrm{rec}}$ encourages the quantized latent to preserve information needed for triple reconstruction, while $\mathcal{L}_{\mathrm{VQ}}$ aligns the continuous encoder outputs with the discrete residual codebook space.

Reconstruction Fidelity Metrics. We report reconstruction recall ($\mathit{RecRecall}$), precision ($\mathit{RecPrecision}$), and their harmonic mean ($\mathit{RF1}$) as diagnostics of how faithfully the learned latent space reconstructs input triples. Detailed definitions are provided in the Appendix.

Observation-Anchored Latent Graph Diffusion

After the knowledge graph autoencoder is trained, we use it to map both observed and target KGs into the latent space. The second stage (Stage II) of LatentDiff is to learn a conditional diffusion model $p_{\theta}(Z_0 \mid Z_{\mathrm{obs}})$, where $Z_{\mathrm{obs}}$ is the latent representation of the observed KG and $Z_0$ is the latent representation of the target KG.

Latent Priors Construction

During Stage-II training, we construct paired observed-target KGs from the training graph. Specifically, we sample $\rho_{\mathrm{tar}} \sim U(\rho_{\min}, 1.0)$ and retain each training triple with probability $\rho_{\mathrm{tar}}$ to construct $G_{\mathrm{tar}}$. Then, we obtain $G_{\mathrm{obs}}$ by retaining each triple in $G_{\mathrm{tar}}$ with probability $\rho_{\mathrm{obs}}$, ensuring $\mathcal{T}_{\mathrm{obs}} \subseteq\mathcal{T}_{\mathrm{tar}} \subseteq\mathcal{T}_{\mathrm{train}}$. With the trained encoder and quantizer in Stage I, we have

where $E_{\phi}$ is the relation-aware graph encoder and $Q_{\phi}$ is the chunked residual quantizer. The autoencoder parameters are frozen during diffusion training.

Mean-Reverting Forward Process

Standard diffusion models perturb data toward an unconditional Gaussian prior. However, for TSP, the observed graph already provides a strong structural anchor. Therefore, we adopt a mean-reverting forward process following prior restoration diffusion models [30] and specialize its endpoint and conditioning structure for TSP. Let $Z_0$ denote the clean latent representation of the target KG and $Z_{\mathrm{obs}}$ denote the latent representation of the observed KG. We define a linear variance schedule $\{\beta_s\}_{s=1}^{S}$, where $S$ is the total number of diffusion steps. Following DDPM, we set

Here, $\bar{\alpha}_s$ controls the signal-to-noise ratio at diffusion step $s$. The mean-reverting forward process is defined as

When $s$ is small, $\bar{\alpha}_s$ is close to 1, and the noisy latent remains close to the target latent. When $s$ approaches $S$, $\bar{\alpha}_s$ becomes close to 0, and the noisy latent is mainly centered around the observed-graph latent with Gaussian noise. Equivalently, by defining the missing-information residual $\Delta_0=Z_0-Z_{\mathrm{obs}}$ and $\Delta_s=Z_s-Z_{\mathrm{obs}}$, Eq. 15 reduces to the standard DDPM marginal $q(\Delta_s\mid\Delta_0)=\mathcal{N}(\sqrt{\bar{\alpha}_s}\Delta_0,(1-\bar{\alpha}_s)I)$. Therefore, observation-anchored diffusion acts as powerful residual denoising in the latent representation. We instantiate mean-reverting diffusion on the proposed relation-aware latent space and couple it with observation conditioning and dual-view denoising for TSP.

Dual-view Latent Graph Denoiser

The denoising network predicts the clean target latent from the noisy latent $Z_s$, the observed-graph latent $Z_{\mathrm{obs}}$, and the diffusion step $s$:

Here, $[\cdot;\cdot]$ denotes feature-wise concatenation, $W_{\mathrm{in}}$ and $W_{\mathrm{out}}$ are learnable input and output projections, respectively, and $\operatorname{DLGD}_L$ denotes a stack of $L$ Dual-view Latent Graph Denoiser (DLGD) blocks. The diffusion step is incorporated through FiLM conditioning, i.e., $\operatorname{FiLM}_s(\cdot)$. To efficiently capture dependencies along both the relation and entity axes, each block performs a dual-view linear attention:

Here, $H$ and $H^{\mathrm{out}}$ are the block input and output, while $H_R$ and $H_N$ are intermediate representations after relation-wise and node-wise linear attention, respectively. LN denotes layer normalization and FFN denotes a feed-forward network. Specifically, $\operatorname{LinAttn}_R$ exchanges information across relations associated with the same head entity, whereas $\operatorname{LinAttn}_N$ aggregates information across head entities under the same relation. Their combination enables efficient dual-view dependency modeling over the unified latent representation.

Diffusion Training Objective

During training, we randomly sample a diffusion step $s \sim\operatorname{Uniform}(\{1,\ldots,S\})$ and Gaussian noise $\epsilon\sim\mathcal{N}(0,I)$. The noisy latent $Z_s$ is constructed by the mean-reverting forward process. The denoiser is trained to predict the clean target latent $Z_0$:

Inference and Triple Set Decoding

At inference, the observed KG $G_{\mathrm{obs}}$ is encoded into $Z_{\mathrm{obs}}$. Since $\bar{\alpha}_S \approx0$, we initialize $Z_S = Z_{\mathrm{obs}} + \epsilon$, where $\epsilon\sim\mathcal{N}(0,I)$. At each reverse step $s=S,\ldots,1$, the denoiser predicts $\widehat{Z}_0=f_\theta(Z_s,Z_{\mathrm{obs}},s)$, followed by DDPM sampling:

where $\widetilde{\beta}_s=\beta_s(1-\bar{\alpha}_{s-1})/(1-\bar{\alpha}_s)$ is the posterior variance, and $\mathbb{I}[s>1]$ removes the stochastic noise at the final reverse step. After the final reverse step, we obtain the generated latent $\widehat{Z}_0$. We decode this representation without an additional quantization step. Since the factorized decoder operates on latent vectors instead of discrete code indices, it maps $\widehat{Z}_0$ to candidate triple scores as:

Triples already contained in $T_{\mathrm{obs}}$ are filtered out. Following the standard TSP evaluation protocol, a single validation-selected score threshold is applied to the remaining candidates to obtain the final predicted missing triple set $\widehat{T}_{\mathrm{miss}}$. The same threshold-selection procedure is applied to all compared methods to ensure a consistent evaluation.

Method Details. Detailed dual-view linear attention and complexity analysis are provided in the Appendix.

Experiments

Experimental Setup

Datasets

We evaluate our method on three common benchmark datasets for TSP: Wiki79k, Wiki143k, and CFamily. Table 1 lists the entity, relation, and triple counts for each dataset. We strictly use the original data partitions and definitions provided by GPHT [50], which guarantees that our results are comparable with existing methods.

Datasets	Entity	Relation	Train	Valid	Test
Wiki79k	7,983	85	57,033	6,337	15,843
Wiki143k	13,928	109	103,415	11,490	28,727
CFamily	2,378	12	16,549	1,839	4,598

Evaluation Metrics

To ensure a fair and comprehensive comparison, we adopt the evaluation metrics proposed for the TSP task in GPHT [50], including Joint Precision ($JPrecision$), Square-root Test Recall ($STRecall$), and TSP Score ($F_{\mathrm{TSP}}$). Because $STRecall$ is the square root of conventional test recall, $F_{\mathrm{TSP}}$ differs from conventional $F_1$. More details are in the Appendix.

Main Results

Table 2 and Table 3 report the main results on all datasets. For LatentDiff, we report the mean and standard deviation over three independent runs. We take $F_{\mathrm{TSP}}$ as the primary comparison metric because it balances the precision of the predicted triple set and the recovery of ground-truth triples. Therefore, we rank different methods mainly according to $F_{\mathrm{TSP}}$, while using $JPrecision$ and $STRecall$ to interpret the precision-recall behavior of each method. We also report the standard $F_1$, reconstructed from the reported mean triple counts. Overall, LatentDiff achieves the best $F_{\mathrm{TSP}}$ on all three datasets, outperforming the strongest baseline by absolute 5.4 percentage points on average.

Table 2: The results on Wiki79k and Wiki143k, where $\mathcal{T}^{\mathrm{POWA+}}_{\mathrm{pred}}$ is the number of correctly predicted test triples, and $\mathcal{T}^{\mathrm{POWA}}_{\mathrm{pred}}$ is the triple set that can be judged under the RS-POWA assumption. The best results are in bold, and the second-best results are underlined. $F_{\mathrm{TSP}}$ is the primary comparison metric. †: Standard $F_1$ under RS-POWA.

Datasets	Models	Number of Triples in RS-POWA			RS-POWA Metrics				CWA Metrics
		$\mathcal{T}_{\mathrm{pred}}$	$\mathcal{T}^{\mathrm{POWA}}_{\mathrm{pred}}$	$\mathcal{T}^{\mathrm{POWA+}}_{\mathrm{pred}}$	$J$Precision	$ST$ Recall	$F_{\mathrm{TSP}}$	$F_1^{\dagger}$	$J$Precision	$F_{\mathrm{TSP}}$
Wiki79k	HAKE [52]	$125\pm87$	$119\pm83$	$30\pm7$	$0.246\pm14.8\%$	$0.044\pm0.4\%$	$0.075\pm0.3\%$	$0.004$	$0.240$	$0.074$
	PairRE [4]	$29742\pm105$	$10217\pm51$	$2901\pm7$	$0.191\pm0.1\%$	$0.428\pm0.0\%$	$0.264\pm0.2\%$	$0.187$	$0.097$	$0.158$
	AstarNet [54]	$5689\pm287$	$1750\pm116$	$1328\pm52$	$0.496\pm0.8\%$	$0.289\pm0.5\%$	$0.365\pm0.2\%$	$0.143$	$0.233$	$0.258$
	ULTRA [13]	$9743\pm0$	$2945\pm0$	$1062\pm0$	$0.235\pm0.0\%$	$0.258\pm0.0\%$	$0.246\pm0.0\%$	$0.104$	$0.109$	$0.153$
	GPHT(HAKE) [50]	$209\pm47$	$191\pm47$	$44\pm8$	$0.220\pm8.0\%$	$0.053\pm0.4\%$	$0.085\pm1.1\%$	$0.006$	$0.210$	$0.084$
	GPHT(PairRE) [50]	$12392\pm3813$	$5866\pm1262$	$2018\pm332$	$0.253\pm1.6\%$	$0.357\pm2.8\%$	$0.296\pm0.4\%$	$0.170$	$0.162$	$0.224$
	SMuRT [48]	$13450\pm0$	$3197\pm0$	$1245\pm0$	$0.241\pm0.0\%$	$0.280\pm0.0\%$	$0.259\pm0.0\%$	$0.118$	$0.093$	$0.139$
	DiffTSP [17]	$7472\pm103$	$4355\pm51$	$3472\pm87$	$0.630\pm1.2\%$	$0.468\pm0.5\%$	$0.537\pm0.5\%$	$0.325$	$0.464$	$0.466$
	LatentDiff	$7384\pm52$	$4871\pm34$	$4060\pm23$	$0.692\pm0.5\%$	$0.506\pm0.1\%$	$0.585\pm0.3\%$	$0.374$	$0.550$	$0.527$
Wiki143k	HAKE [52]	$22215\pm1283$	$6182\pm43$	$3044\pm72$	$0.315\pm0.2\%$	$0.326\pm0.3\%$	$0.32\pm0.3\%$	$0.159$	$0.137$	$0.193$
	PairRE [4]	$19228\pm2075$	$3313\pm722$	$1191\pm77$	$0.211\pm3.0\%$	$0.204\pm0.7\%$	$0.207\pm1.0\%$	$0.070$	$0.061$	$0.095$
	AstarNet [54]	$6001\pm231$	$2589\pm129$	$2275\pm31$	$0.629\pm0.8\%$	$0.281\pm0.2\%$	$0.389\pm0.2\%$	$0.140$	$0.379$	$0.323$
	ULTRA [13]	$3844\pm0$	$1349\pm0$	$692\pm0$	$0.346\pm0.0\%$	$0.155\pm0.0\%$	$0.214\pm0.0\%$	$0.045$	$0.180$	$0.166$
	GPHT(HAKE) [50]	$17702\pm7935$	$4709\pm60$	$2700\pm681$	$0.363\pm4.8\%$	$0.307\pm4.2\%$	$0.333\pm4.6\%$	$0.150$	$0.152$	$0.204$
	GPHT(PairRE) [50]	$3011\pm233$	$1954\pm191$	$909\pm34$	$0.384\pm1.9\%$	$0.178\pm0.3\%$	$0.243\pm0.1\%$	$0.059$	$0.302$	$0.224$
	DiffTSP [17]	$10162\pm530$	$6804\pm275$	$4543\pm105$	$0.557\pm0.3\%$	$0.397\pm0.4\%$	$0.464\pm0.6\%$	$0.246$	$0.447$	$0.420$
	LatentDiff	$11211\pm57$	$5748\pm20$	$5149\pm42$	$0.677\pm0.5\%$	$0.423\pm0.2\%$	$0.521\pm0.1\%$	$0.283$	$0.459$	$0.440$

Table 3: The results on the CFamily dataset, where $\mathcal{T}^{\mathrm{CWA+}}_{\mathrm{pred}}$ is the set of correctly predicted test triples. The best results are in bold, and the second-best results are underlined. $F_{\mathrm{TSP}}$ is the primary comparison metric. †: Standard $F_1$ under CWA.

Models	Number of Triples in			CWA Metrics
	$\mathcal{T}_{\mathrm{pred}}$	$\mathcal{T}^{\mathrm{CWA}}_{\mathrm{pred}}$	$\mathcal{T}^{\mathrm{CWA+}}_{\mathrm{pred}}$	$J$Precision	$ST$ Recall	$F_{\mathrm{TSP}}$	$F_1^{\dagger}$
HAKE [52]	$3186\pm543$	$3186\pm543$	$1788\pm321$	$0.561\pm1.8\%$	$0.624\pm5.9\%$	$0.591\pm3.1\%$	$0.459$
PairRE [4]	$5732\pm2062$	$5732\pm2062$	$1747\pm694$	$0.305\pm1.8\%$	$0.616\pm13.7\%$	$0.408\pm4.9\%$	$0.338$
RGCN [37]	$30608\pm3443$	$30608\pm3443$	$3026\pm137$	$0.099\pm0.6\%$	$0.811\pm0.5\%$	$0.176\pm5.2\%$	$0.172$
CompGCN [41]	$38931\pm3745$	$38931\pm3745$	$2599\pm47$	$0.067\pm4.8\%$	$0.752\pm0.3\%$	$0.123\pm1.7\%$	$0.119$
AstarNet [54]	$2075\pm360$	$2075\pm360$	$911\pm187$	$0.439\pm1.7\%$	$0.445\pm4.6\%$	$0.442\pm3.3\%$	$0.273$
ULTRA [13]	$32761\pm0$	$32761\pm0$	$763\pm0$	$0.023\pm0.0\%$	$0.407\pm0.0\%$	$0.044\pm0.0\%$	$0.041$
GPHT(HAKE) [50]	$1896\pm149$	$1896\pm149$	$1222\pm58$	$0.645\pm2.6\%$	$0.516\pm1.2\%$	$0.573\pm0.4\%$	$0.376$
GPHT(PairRE) [50]	$3739\pm593$	$3739\pm593$	$1471\pm187$	$0.393\pm1.3\%$	$0.566\pm3.4\%$	$0.464\pm0.7\%$	$0.353$
LLM TSP(GPT-3.5) [49]	$3403\pm158$	$3403\pm158$	$96\pm12$	$0.028\pm0.3\%$	$0.145\pm0.8\%$	$0.047\pm0.4\%$	$0.024$
LLM TSP(GPT-4o) [49]	$1276\pm147$	$1276\pm147$	$179\pm17$	$0.140\pm0.5\%$	$0.197\pm0.9\%$	$0.164\pm0.3\%$	$0.061$
SMuRT [48]	$2358\pm0$	$2358\pm0$	$657\pm0$	$0.279\pm0.0\%$	$0.378\pm0.0\%$	$0.321\pm0.0\%$	$0.189$
DiffTSP [17]	$2453\pm35$	$2453\pm35$	$1657\pm18$	$0.675\pm0.2\%$	$0.600\pm3.0\%$	$0.635\pm0.2\%$	$0.470$
LatentDiff	$3299\pm60$	$3299\pm60$	$2252\pm39$	$0.683\pm0.5\%$	$0.700\pm0.5\%$	$0.691\pm0.2\%$	$0.570$

Ablation Study

Table 4 reports the ablation results on three datasets using $F_{\mathrm{TSP}}$ as the primary metric. The full LatentDiff achieves the best performance on all datasets, demonstrating the effectiveness of the overall design. LatentDiff-Subgraph retains the same modules as LatentDiff, but partitions the KG into local subgraphs like GPHT. The performance drop provides a controlled comparison that supports the benefit of operating in a unified graph latent space. For Stage I, w/o CRVQ replaces CRVQ with a codebook parameter matched VQ-VAE while keeping the encoder, decoder, latent dimensionality, training protocol, and total number of learnable codebook parameters comparable. Its substantial degradation is therefore not attributable to a smaller parameter budget. VQ-VAE exhibits poor Stage-I reconstruction fidelity and provides an inadequate latent space for subsequent diffusion. Detailed results are provided in the Appendix. w/o Quantization removes the quantization module and directly uses the continuous encoder output as the latent representation. The degradation of w/o Quantization suggests that unconstrained continuous latents are less suitable for diffusion-based TSP, supporting the role of quantization in constructing a structured latent space. Removing the relational-view attention (w/o Relational view Attention) or replacing the factorized decoder (w/o Factorized Decoder) with MLPs also leads to clear performance drops. w/o Stage II directly decodes the observed latent. Its poor performance shows

Variant	Wiki79k	Wiki143k	CFamily	Avg.
LatentDiff	0.585	0.521	0.691	0.599
LatentDiff-Subgraph	0.539	0.472	0.653	0.555
w/o Relational view Attention	0.576	0.495	0.674	0.582
w/o CRVQ (VQ-VAE)	0.030	0.075	0.455	0.187
w/o Quantization (Vanilla AE)	0.498	0.414	0.332	0.415
w/o Factorized Decoder	0.557	0.481	0.659	0.566
w/o Stage II (Direct Decoding)	0.277	0.221	0.351	0.283
w/o Iterative Diffusion	0.458	0.442	0.518	0.473
Deterministic Iterative Refinement	0.501	0.471	0.615	0.529
Standard Gaussian Diffusion	0.513	0.476	0.637	0.542
w/o Relation-wise Linear Attention	0.538	0.493	0.554	0.528
w/o Node-wise Linear Attention	0.546	0.508	0.602	0.552

that our proposed Stage-II module provides the conditional latent transformation required for TSP. For Stage II, standard Gaussian diffusion decreases the average score to 0.542, confirming the importance of the mean-reverting diffusion process that anchors generation on the observed KG. w/o Iterative Diffusion replaces diffusion with a one-step predictor while retaining the same denoiser backbone, the resulting performance drop demonstrates the benefit of progressive latent diffusion refinement. We further compare LatentDiff with a deterministic iterative refinement baseline using the same denoiser backbone and refinement steps. Unlike diffusion, this baseline learns to progressively update the latent representation without a noise-based forward process. Its lower performance suggests that the benefit of LatentDiff extends beyond repeated latent refinement. In addition, removing either relation-wise (w/o Relation-wise Linear Attention) or node-wise linear attention (w/o Node-wise Linear Attention) consistently hurts performance, suggesting that both relation-level and node-level dependencies contribute to effective latent generation.

Construction of LatentDiff-Subgraph. We use exactly the graph partitioning procedure and partition hyperparameters of GPHT for each dataset. The same LatentDiff parameters are shared across all subgraphs. Predictions produced by multiple overlapping subgraphs are merged by union, while triples whose endpoints never co-occur in a subgraph cannot be generated. The prediction threshold is independently selected on the validation set for this variant. All hyperparameters are re-tuned using the same search space and validation budget as LatentDiff.

Hyperparameter Sensitivity Analysis

Effect of Codebook Size

We study the effect of codebook size in the CRVQ module. As shown in Figure 3, increasing the codebook size consistently improves the reconstruction quality in Stage I. However, the TSP performance saturates when the codebook size reaches 1024.

Effect of CRVQ Configuration

We analyze the two key hyperparameters in CRVQ: the number of chunks and the number of residual quantization levels. As shown in Figure 4, increasing the number of chunks substantially improves reconstruction quality, while a moderate number of residual codebooks further benefits TSP performance.

Effect of Diffusion Steps

We further investigate the effect of the number of diffusion steps. As shown in Figure 5, using only 2 diffusion steps leads to clearly inferior performance. When the number of steps increases from 2 to 10 and then to 50, the $F_{\mathrm{TSP}}$ score improves substantially. However, the performance saturates after 50 or 100 steps, and further increasing the steps to 200 does not yield consistent gains.

Case Study

We provide a qualitative case study to illustrate the failure mode of subgraph-based TSP baselines. A cross-subgraph triple refers to a test triple whose head and tail entities do not appear in any common subgraph of subgraph-based methods. Table 6 shows representative cross-subgraph test triples recalled by LatentDiff using the same validation-selected threshold as in the main evaluation. The compared subgraph-based methods do not generate these triples.

Target Triple	GPHT	DiffTSP	LatentDiff
(October, follows, September)	Missed	Missed	Recalled ($s = 0.98$)
(2006 Asian Games, follows, 2002 Asian Games)	Missed	Missed	Recalled ($s = 0.97$)
(Châlons-en-Champagne, twinned with, Neuss)	Missed	Missed	Recalled ($s = 0.95$)
(Épinal, located in, Vosges)	Missed	Missed	Recalled ($s = 0.91$)
(Braga, shares border with, Guimarães)	Missed	Missed	Recalled ($s = 0.89$)
(Bonn, country, Weimar Republic)	Missed	Missed	Recalled ($s = 0.75$)

Analysis of Cross Subgraph Recovery

A key motivation of our framework is that subgraph partitioning may fragment relational dependencies. To directly examine this issue, we construct a cross-subgraph test subset based on the partitions used by the subgraph-based baseline. A test triple $(h,r,t)$ is considered cross-subgraph if its head entity $h$ and tail entity $t$ are not contained in any common subgraph. Such triples are difficult for subgraph-based methods because the evidence associated with $h$ and $t$ is separated by the partition boundary.

We examine how much structural information is affected by the partition process. Given the subgraph collection used by subgraph-based baseline, we define a training triple as covered if it appears in at least one subgraph. Otherwise, this triple is counted as partition-uncovered. Similarly, a test triple is defined as cross-subgraph if its head and tail entities do not appear in any common subgraph. Table 5 reports the statistics. On Wiki79k, 5.43% training triples are uncovered after partitioning, and 21.40% of test triples are partition-boundary triples. The effect is more pronounced on Wiki143k. CFamily shows smaller fragmentation, with 1.02% partition-boundary test triples. These results suggest that the impact of subgraph partitioning varies across datasets, but can be non-negligible on larger KGs. This motivates our unified graph latent formulation, where all head-relation latent embeddings are represented in a globally shared code latent space.

Dataset	#Train	#Retained	#Uncovered	Uncovered Ratio	#Test	#Same-subgraph	#Partition-boundary	Boundary Ratio
CFamily	16,549	16,514	35	0.21%	4,598	4,551	47	1.02%
Wiki79k	57,033	53,937	3,096	5.43%	15,843	12,452	3,391	21.40%
Wiki143k	103,415	88,984	14,431	13.95%	28,727	16,781	11,946	41.58%

Performance on partition-boundary triples. We further conduct a subset-level evaluation on partition-boundary test triples, whose head and tail entities do not occur in any common subgraph constructed by the subgraph-based models. As shown in Table 7, GPHT and DiffTSP recover none of these triples. In contrast, LatentDiff can score the partition-boundary triples, resulting in $STRecall_{bd}$ values of 0.229 and 0.195 on Wiki79k and Wiki143k, respectively. Partition-boundary triples remain difficult because their supporting evidence can be sparse or structurally distant, and the decoding threshold is selected according to overall validation $F_{TSP}$ rather than boundary-specific performance. Nevertheless, the non-zero recovery of LatentDiff, compared with zero recovery by the subgraph-based methods, provides quantitative evidence that the unified latent generation framework enables the recovery of triples excluded by subgraph-restricted candidate construction.

Dataset	Method	Recalled	$STRecall_{bd}$
Wiki79k	GPHT	0/3,391	0.000
Wiki79k	DiffTSP	0/3,391	0.000
Wiki79k	LatentDiff	178/3,391	0.229
Wiki143k	GPHT	0/11,946	0.000
Wiki143k	DiffTSP	0/11,946	0.000
Wiki143k	LatentDiff	456/11,946	0.195

Efficiency Study

We further compare LatentDiff with DiffTSP and GPHT under the same hardware environment. As shown in Table 8, LatentDiff achieves faster inference. These results show that performing generation in the latent space can provide favorable inference efficiency despite the iterative nature of diffusion sampling. LatentDiff also avoids explicitly constructing the raw $|\mathcal{E}| \times|\mathcal{R}| \times|\mathcal{E}|$ adjacency tensor and instead performs generation over $\mathbf{Z} \in\mathbb{R}^{|\mathcal{E}| \times|\mathcal{R}| \times d}$. Moreover,

Dataset	GPHT	DiffTSP	LatentDiff	Speedup
CFamily	46.25	8.13	1.64	4.9×
Wiki79k	151.94	83.77	12.79	6.5×
Wiki143k	583.69	101.07	35.93	2.8×

head-wise chunking reduces peak latent memory without partitioning the graph or introducing separate local model parameters. More details are in the Appendix.

More Experimental Details. Additional analysis on ablation studies, implementation details, observed ratios, and inference memory consumption are reported in the Appendix.

Conclusion

In this work, we proposed LatentDiff, a two-stage conditional latent generative framework, for knowledge graph triple set prediction. LatentDiff first learns a relation-aware residual-quantized latent space, and then performs observation-anchored diffusion in this space to recover missing relational residuals. Instead of relying on subgraph-level local prediction, LatentDiff performs completion in a globally shared latent space. Experimental results demonstrate that LatentDiff achieves state-of-the-art performance. This work focuses on establishing the effectiveness of unified latent generation for TSP, and a systematic investigation of its scalability to larger knowledge graphs is left for future work.

Method Details

Dual-view Linear Attention

Given a sequence $X \in\mathbb{R}^{L \times d}$, each attention head computes

and uses the positive feature map $\phi(x) = \operatorname{ELU}(x) + 1$:

We employ multi-head attention and concatenate the outputs of all heads. Relation-wise attention applies this operation over the $R$ relation tokens for each entity, whereas node-wise attention applies it over the $N$ entity tokens for each relation. The two branches use independent projection parameters.

Complexity Analysis

Let $N = |\mathcal{E}|$, $R = |\mathcal{R}|$, and $d$ denote the entity count, relation count, and latent dimension. Let $L_e$, $L_d$, and $L_c$ denote the encoder, denoiser, and residual quantization depths, $C$ the codebook size, and $S$ the number of reverse diffusion steps. Raw adjacency space reconstruction or diffusion requires $O(N^2R)$ memory and at least $O(N^2Rd)$ computation. Instead, LatentDiff generates head-relation latents $Z \in\mathbb{R}^{N \times R \times d}$ with $O(NRd)$ latent memory. Stage I. For input triples $\mathcal{T}_{\rho}$, sparse projection costs $O(|\mathcal{T}_{\rho}|d)$, while reconstruction with $K_{\mathrm{neg}}$ negatives per positive costs $O((1 + K_{\mathrm{neg}})|\mathcal{T}_{\rho}|d)$. Relation-wise encoding and CRVQ respectively cost

CRVQ uses $O(L_cC d)$ codebook memory. Splitting embeddings into proportionally smaller chunks leaves its leading order computation and parameter memory unchanged.

Stage II. Dual-view linear attention avoids dense $R \times R$ and $N \times N$ attention maps and their $O(NR^2d + R^2Nd)$ computation. With attention feature dimension and FFN width both $O(d)$, each denoiser evaluation costs $O(L_dNRd^2)$. Training samples one timestep and has the same per-instance asymptotic cost, whereas $S$-step DDPM sampling costs $O(SL_dNRd^2)$. Reverse steps execute sequentially without accumulating activations across steps, so peak memory does not scale with $S$.

Chunking and decoding. Scoring a candidate set $\mathcal{C}$ costs $O(|\mathcal{C}|d)$, or $O(N^2Rd)$ for exhaustive decoding, while scores can be generated and thresholded incrementally. In both stages, processing $B_h$ head entities per chunk reduces active latent memory and exhaustive score buffer memory, respectively, as

Total asymptotic computation remains unchanged. Chunking manages memory rather than partitioning the graph: all chunks share the encoder, codebooks, denoiser, and decoder within a common latent space. Global node-wise interactions can be preserved by accumulating key-value statistics across chunks rather than computing attention independently within each chunk.

Stage-I Reconstruction Fidelity Metrics

Stage I learns a compact latent representation by reconstructing its input graph. Its reconstruction metrics measure how faithfully an input is encoded and decoded. The diagnostic evaluation below supplies an evaluation graph containing test triples to the encoder, so its scores describe reconstruction with access to those triples. These diagnostic inputs and scores are excluded from TSP training and hyperparameter selection. Stage-I reconstruction evaluation is performed by feeding an evaluation KG containing both training and test triples:

After encoding, quantization, and decoding, we obtain a reconstructed triple set $\mathcal{T}_{\mathrm{rec}}$. Since training triples are already trained in our model and should not inflate the reconstruction score, we remove them from the reconstructed set:

The reconstruction metrics are then computed only against the test triples:

Here, $RF_1$ is used as the primary diagnostic metric for Stage-I reconstruction analysis.

Experiment Details

Baselines

We select both link prediction methods and TSP methods as baselines. The representative link prediction methods, including HAKE [52], PairRE [4], RGCN [37], CompGCN [41], AstarNet [54] and ULTRA [13], are modified for TSP by scoring all candidate triples in the graph and selecting those above a threshold as the predicted results. The threshold is determined through grid search on the validation set to achieve the best performance. The latest TSP methods, such as GPHT [50], LLMTSP [49], SMuRT [48] and DiffTSP [17], are also included.

Evaluation Metrics

TSP Evaluation Metrics. $\mathit{JPrecision}$ measures the precision of the predicted triple set while accounting for the size of the set itself. $\mathit{STRecall}$ measures the percentage of test triples successfully recovered by the model with a square-root transformation. $F_{\mathrm{TSP}}$ is the harmonic mean of $\mathit{JPrecision}$ and $\mathit{STRecall}$. Higher values indicate better performance for all metrics.

Following GPHT, we denote the predicted missing triple set as $\mathcal{T}_{\mathrm{pred}}$. This corresponds to $\mathcal{T}_{\mathrm{bmiss}}$ in our problem definition. Under the closed-world assumption (CWA), the predicted triples are divided into correct and incorrect parts:

Under the relation-similarity partial open-world assumption (RS-POWA), the positive predictions are defined as

The negative part is defined as

Then the valid predicted set under RS-POWA is

For $WA\in\{\mathrm{CWA},\mathrm{POWA}\}$, the three metrics are computed as

Since $F_{\mathrm{TSP}}$ is the harmonic mean of $\mathit{JPrecision}$ and $\mathit{STRecall}$, it summarizes the trade-off between set precision and test triple recovery. Therefore, we use $F_{\mathrm{TSP}}$ as the primary metric for model comparison, following the TSP evaluation protocol in GPHT. $\mathit{JPrecision}$ and $\mathit{STRecall}$ are reported as complementary metrics to analyze whether the improvement comes from more accurate predicted sets or better recovery of test triples. Unless otherwise specified, our discussion mainly focuses on $F_{\mathrm{TSP}}$. We report micro-averaged results over all predicted triples.

Triple Set Decoding and Threshold Selection

For each method, the final predicted triple set is obtained by thresholding its predicted triple scores. Let $s_m(h,r,t)$ denote the score assigned by method $m$ to a candidate triple $(h,r,t)$. Given a threshold $\tau$, the predicted triple set is

where $\mathcal{C}$ denotes the candidate triple set after filtering out observed training triples. Following the evaluation protocol used in GPHT and DiffTSP, we select the threshold on the validation set for each method and each dataset. Specifically, we search over a predefined threshold grid $\Theta$ and choose

The selected threshold $\tau_m^*$ is then fixed and used to generate the predicted triple set on the test set. All reported $\mathit{JPrecision}$, $\mathit{STRecall}$, and $F_{\mathrm{TSP}}$ values are computed on the test set using this validation-selected threshold. This protocol is applied consistently to LatentDiff and all compared baselines to ensure a fair comparison.

Implementation Details

All experiments are conducted on one NVIDIA A100 GPU. We use the Adam optimizer. Hyperparameters are selected by grid search on the validation set, more details are in Table 9. To further reduce peak memory usage, both Stage I and Stage II support head-wise chunking implementation. Specifically, the latent tensor is processed in chunks along the head-entity dimension instead of materializing all head-relation embeddings at once. This is an engineering strategy for memory-efficient training and inference, not a subgraph partitioning strategy.

Hyperparameter	CFamily	Wiki79k	Wiki143k
Stage-I Training
Encoder hidden dimension	64	64	32
Chunks of CRVQ $K$	4	4	4
Number of residual quantizers $L_c$	4	4	2
Latent dimension	16	16	16
Codebook size $C$	1024	1024	2048
Commitment loss weight $\lambda_m$	1	1	1
CRVQ loss weight $\lambda_{\mathrm{VQ}}$	0.25	0.25	0.25
Number of negative samples $K_{\mathrm{neg}}$	512	256	256
Learning rate	$5 \times10^{-4}$	$5 \times10^{-4}$	$5 \times10^{-4}$
Stage-II Training
Ratio $\rho_{\min}$	0.7	0.7	0.7
Observed KG ratio $\rho_{\mathrm{obs}}$	0.8	0.8	0.8
Variance schedule	Linear	Linear	Linear
Variance schedule $\beta_{\min}$	$1 \times10^{-3}$	$2 \times10^{-3}$	$2 \times10^{-3}$
Variance schedule $\beta_{\max}$	$2 \times10^{-1}$	$4 \times10^{-1}$	$4 \times10^{-1}$
Diffusion steps	100	50	50
Denoising hidden dimension	128	128	64
Condition dimension	128	128	64
Time embedding dimension	128	128	64
Denoiser depth $L$	3	2	2
Learning rate	$5 \times10^{-4}$	$5 \times10^{-4}$	$5 \times10^{-4}$

Ablation Study Details

Table 10 reports the Stage-I reconstruction ablations. For the VQ-VAE comparison, we match the total learnable codebook parameter budget while keeping the encoder, decoder, latent dimensionality, and training protocol identical to those of LatentDiff. Replacing CRVQ with VQ-VAE causes the largest drop in reconstruction fidelity, supporting the effectiveness of CRVQ in preserving KG structure. Removing relational-view attention or replacing the factorized decoder also reduces reconstruction quality, highlighting the contributions of relation-aware encoding and factorized decoding.

Variant	CFamily	Wiki79k	Wiki143k
LatentDiff	0.9703	0.6918	0.7145
w/o Relational view Attention	0.9057	0.6229	0.6927
w/o CRVQ (VQ-VAE)	0.4591	0.0575	0.1098
w/o Quantization (Vanilla AE)	0.9898	0.7187	0.7339
w/o Factorized Decoder	0.8895	0.5891	0.6219

Although the Vanilla AE achieves slightly higher reconstruction fidelity, its TSP performance is substantially worse (Table 4). This contrast suggests that reconstruction fidelity alone is insufficient for effective latent generation and is consistent with a regularizing role of quantization. Stage-I reconstruction fidelity is a diagnostic of latent representation quality rather than a strict upper bound on TSP performance, since Stage II directly decodes continuous latent predictions without re-quantization.

Sensitivity to the Observed Ratio

We further study the effect of the observed ratio, i.e., the proportion of triples retained as the observed graph for conditioning in the diffusion training process. Figure 6 shows that the performance is relatively stable across different observed ratios, but the best results are achieved when the observed ratio is around 0.7–0.8. When the observed ratio becomes too high, the prediction task becomes less informative because fewer target triples remain to be recovered, and the model may also receive reduced supervision on the missing part. These results suggest that a moderate observed ratio provides a good balance between observation information and prediction difficulty. Therefore, we use an observed ratio of 0.8 as the default setting.

Effect of CRVQ Configurations

As shown in Figure 7, we also analyze two key hyperparameters of CRVQ on Wiki79k. Increasing either the number of chunks or the number of residual codebooks improves both Stage-I reconstruction fidelity and TSP performance.

Inference Efficiency and Memory Cost

We evaluate the inference efficiency of LatentDiff on one A100 GPU. Table 11 reports the prediction time and peak memory under different diffusion steps. The inference time increases with $S$, since

Dataset	Prediction Time (s)		Peak Memory (MB)
	$S = 50$	$S = 100$	$S = 50$	$S = 100$
CFamily	0.72	1.64	2,136	2,136
Wiki79k	12.79	22.96	12,366	12,366
Wiki143k	35.93	45.76	19,872	19,872

DDPM sampling requires iterative reverse denoising. In contrast, the peak memory remains unchanged when increasing $S$, because the reverse process is executed sequentially and does not store activations from all denoising steps.