Cycle Conditioning for Robust Representation Learning from Categorical Data
Authors: Mohsen Tabejamaat
Farzaneh Etminani
Mattias Ohlsson
Venue: Accepted by TMLR ()

Abstract: This paper introduces a novel diffusion-based method for learning representations from categorical data. Conditional diffusion models have demonstrated their potential to extract meaningful representations from input samples. However, they often struggle to yield versatile, general-purpose information, limiting their adaptability to unforeseen tasks. To address this, we propose a cycle conditioning approach for diffusion models, designed to capture expressive information from conditioning samples. However, cycle conditioning alone can be insufficient. Diffusion models may ignore conditioning samples that vary across training iterations, an issue that occurs within cycle conditioning. To counter this limitation, we introduce additional "spelling" information to guide the conditioning process, ensuring that the conditioning sample remains influential during denoising. While this supervision enhances the generalizability of extracted representations, it is constrained by the sparse nature of spelling information in categorical data, leading to sparse latent conditions. This sparsity reduces the robustness of the extracted representations for downstream tasks or as effective guidance in the diffusion process. To overcome this challenge, we propose a linear navigation strategy within the latent space of conditioning samples, allowing dense representations to be extracted even with sparse supervision. Our experiments demonstrate that our method achieves at least a 1.42\% improvement in AUROC and a 4.12\% improvement in AUCPR over the best results from existing state-of-the-art methods.

Published in Transactions on Machine Learning Research (03/2025)
Cycle Conditioning for Robust Representation Learning from Categorical Data
Mohsen Tabejamaat
mohksen.tabejmmaathhhhh
School of Information Technology, Halmstad University, Sweden
Farzanch Etminani
Centrum for Research and Innovation, Region Hallond, Sweden School of Information Technology, Halmstad University, Sweden
Mattias Ohlsson
Colliias.ohlsson@@cce
Centre for Environmental and Climate Science, Land University, Sweden School Information Technology, Halmstad University, Sweden
Reviewed on OpenReview: https: openreviee. net/ forum? GkYOcbNLaW
Abstract
This paper introduces novel diffusion-basse method for learning representations from categorical data. Conditional diffuusio models have demonstrated their potential to extract meaningful representations from input samples. However, they often struggle to yield versatile, multi-lask information, limiting their adaptability to unforeseen tasks. To address this, we propose a cycle conditioning approach for diffusion models, designed to capture expressive information from conditioning samples. However, cycle conditioning alone can be insufficient. Diffusion models may ignore conditioning samples that vary across training iterations, an issue that occurs within cycle conditioning. To counter this limitation, we introduce additional "spelling' information to guide the conditioning process, ensuring that the conditioning sample remains influential during denoising. While this supervision enhances the generalizability of extracted representations, it constrained by the sparse nature of spelling information in categorical data, leading sparse latent conditions. This sparsity reduces the robustness of the extracted representations for downstream tasks effective guidance in the diffusion process. To overcome this challenge, propose linear navigation strategy within the latent space of conditioning samples, allowing dense representations to be extracted even with sparse supervision. Our experiments demonstrate that our method achieves at least a 1.42% improvement AUROC and a 4.12% improvement in AUCPR over the best results from existing state-of-the-art methods.
1 Introduction
A categorical time series is a sequence of observations where each observation corresponds to specific category rather than a numerical value, making interpolation between observations neither possible nor meaningful. Healthcare trajectories are prominent example, often consisting of large categorical time series with observations representing medical activities, like prescribed medications treatment procedures. Effectively utilizing such sequences requires creating compact and mcaningful representations of these observations through process called representation learning. This step is crucial for improving the accuracy and computational efficiency of downstream tasks. While unsupervised representation learning has proven effective for continuous data (Lo-Khac et al. 2020 Poudel al.] 2024 Jang & Wang 2023; Izacard al. 2021] Caron et al.] (220)), the need for various augmentations has made difficult to benefit from this unsupervised strategy in learning representation for categorical sequences. For categorical data, augmentation

---

Published in Transactions on Machine Learning Research (03/2025)
is mostly limited to token masking which helps with missing tokens (part or whole of an observation) but fails to capture any complex relationships between the tokens. This limitation reduces the effectiveness of unsupervised representation learning for categorical sequences.
Utilizing representations learned within the latent space of generative models offers a promising alternative that eliminates the need for data augmentation in representation learning tasks (Afkanpour et al. 2024). Tn this approach, a generative model is first trained accurately capture the underlying distribution of the data. The representation of imput samples is then derived from their projection onto the bottleneck layer of the generative model (Abdal et al., 2019: Radford et al., 2015 Zhang et al. 2022a). Recently, diffusion models have demonstrated significant effectivenoss in achioving high-quality fits to data distributions, which has led researchers to investigale their potential in representation learning (Yang et al.) 2024; Hudson et al.) 2024 Preechakul et al., 2022). However, as these models do not include a bottleneck layer, attention has shifted to conditional diffusion models (Yue et al. 2024). These models allow for projecting a conditioning sample into a latent space, which can then be used to guide the denoising process of the diffusion model. While conditioning effectively directs the generative process toward specifie classes of samples, it does not. inherently guarantee the extraction of versatile information from the conditioning samples. Consequently, the model may overlook essential, multi-task information within the conditioning samples that is necessary for good performance of these representations in an unforeseen downstream task.
In addition, these methods have primarily been developed for continuous visual data, which aligns with the continuous nature of diffusion models. The forward process of diffusion models operates on the assumption that adding Gaussian noise to a noisy sample produces a result that can be sampled from a Gaussian distribution. This implies that operations in the latent space of diffusion models occur within a continuous framework. As a result, these methods cannot be directly applied to representation learning from categorical datasets. Although there have been attempts to adapt diffusion processes for discrete data distributions [Hoogeboom et al. 2021 Austin et al., 2021), extending these methods to incorporate a conditioning sample romains challonging, particularly when the condition is derived from the output layer of an encoder, which is inherently contimuous. This difficulty limits the applicability of these approaches in lasks such as represent a ation learning.
In this paper, we show that cycle conditioning in diffusion models can effectively extract multi-task repre- sentations from calegorical time series. To address the mismatch between ca categorical data and continuous diffusion models, we propose a learnable linear transformation that converts each categorical series into a continuous representation. This transformation does not require feedback from the representation learning task, allowing it to be pretrained independently before the initiation of the representation learning process. We also introduce a new strategy to encode the order of numbers and letters in each token of a time series, called spelling structure. This spelling structure used to supervised the learned representations that are further used as conditions for diffusion model. This supervision ensures that diffusion process does not overlook its condition that due to the cycle conditioning changes in its training iteration. Therefore, the network learn to effectively include the representation of conditioning samples in the reconstruction task of the diffusion model. Our strategy ensures that both the order of letters and the contextual relationships between tokens are effectively used in the representation learning process through self-supervised learning. As a result, the learned representations are versatile, making them applicable to any downstream task without concern that they are biased toward specific labels or annotations used during training.
While supervising the latent condition with spelling information aids in creating multi-task representation, it has significant limitations due to the inherent sparsity of spelling information in categorical data, Each token's 's spelling is represented by an index, and the spelling information for an entire time series is encoded as a multi-hot vector with a dimensionality of O", where O is the number of possible characters in the dataset, and n is the length of the longest token. This results in a highly sparse representation. Consequently, using such a sparse vector to supervise the latent condition often produces a sparse latent representation, which is not ideal for capturing multi-task features and lacks the expressiveness needed to provide robust guidance for the diffusion process. To mitigate this issue, we propose achieve this supervision through learnable set of latent codes, where each code represents an initial direction toward a specific element of the spelling
"takens" refer observations their subcomponents, analogous to how large language models handle text.
2

---

Published in Transactions on Machine Learning Research (03/2025)
vector. This setup enables the network to transfer its latent representation along O" meaningful directions before supervising it with the sparse spelling vector. Since these initial directions are learnable and shared across all training samples, the network learns to explore meaningful semantic paths in the latent space, grouping time series that share token spelling similaritics. For instance, two series x = {'A01', ' 'B01'D } and y = ('B01', 'A02' will be closer lo each other than x = {'A01', 'B01') and y = ( {'A01', 'B03'}, }, despite the latter pair having more similarity in token order. This method is crucial for analyzing similarities between irregular time series, where data collection timelines are inconsistent, which is a common scenario in medical applications.
We focus on applying our method to ICD codes that represent medical trajectories and evaluate its performance on the widely recognized MIMIC-III and MIMIC-IV databases. This evaluation highlights the potential of our approach in advancing the understanding of medical data and its applications in clinical practice.
2 Related work
Categorical data representation Early research primarily focused on using metric learning to embed
categorical data into a low-dimensional space (Jian et al., 2018: Li et al., 2023; Zhu al., 2018 Jia et al, 2015). While demonstrating reasonable performance with restricted and balanced mumber of tokens, their effectiveness diminishes in scenarios when there is a wide variation and imbalanced distribution of tokens. Following the successful introduction of mask modeling by BERT (Devlin ct al. 2018),01), several algorithms (Yin el al. 2023 Majmundar et. al., 2022: Du et al. 2023 Yan et. al. 2024: Xie et al. 2024) have emerged aiming to take advantage of this model for representation learning of categorical data. While these methods perform well at capturing meaningful information, they often struggle with handling multidimensional time series data. In another line of research, authors proposed to utilize Graph Neural Networks for categorical representation learning (Liao & Li. 2023; You et al., 2020 Chen et al., 2024ab and demonstrated their effectiveness in capturing intricate dependencies of data. However, the computational expense of these networks hinders their scalability to largo-scale datasets. Contrastive learning between different subsets of categories is another line of research that is considered for categorical data representation (Ucar el al. 2021p. However, the performance of this method heavily depends on the split of the subsets, which is a hyperparameter, requiring additional effort for tuning.
Diffusion models for categorical data There are two approaches for adapting diffusion models to handle categorical data. In the first approach (Dieleman et al. 2022 Lin et al. 2023 Gao et al. 2022). input data is first converted into a continuous domain, and then continuous diffusion model is employed to generate the data from a Gaussian distribution. These methods typically involve training a learnable embedding layer alongside the diffuuson model. However, this setup leads to mutual dependence between the diffusion process and the embedding operation, lcading to a suboptimal solution for both of the networks. Another direction of research explores a discrete diffusson process built, on categorical distributions, thereby enabling their applicability for modeling inputs with binary Heements(Hoogeboodd et al. 2021 Austin et al. [2021] Savinov ct al. 2021 Reid et al., 2022 Inoue et al. 2023). However, these methods have significant difficulties when employed in a conditional mode. This is because the condition, typically provided through conditioning encoder, is continuous, whereas the diffussoo process is discrete.
Representation learning based on diffusion models Learning representations through generative models has become a prominent focus in machine learning research. Much of this research revolves around Generative Adversarial Networks (GANs), (Radford et al. 20155: Chen et al. 2016) and Variational Autoencoders (VAFs) (Wang al. 2023a), both of which benefit from a bottleneck that can be used as representation layer. Recently, conditinnal diffusion models have demonstrated significant potential in discovering rich semantic information from visual data (Mittal et al.]
2023
Preechakul et al., 2022 Zhang et al. 2022b Yang &
Wang. 2023 Wang et al. 2023b Yue et al.)
2024 Our work build on this recent advances in diffusion models for representation learning, notably extending this framework to handle categorical data. However, these methods commonly suffer from a significant drawback: they lack the ability to differentiate between the context that is intended to be contributed by noise and the one that needs to be contributed by the conditioning sample. Consequently, there is limited control over the meaning of the information captured by the representation layer. In contrast, our approach employs a self-supervised strategy that enables to
3

---

Published in Transactions on Machine Learning Research (03/2025)
enrich the representation with both contextual and semantic information of tokens, thereby enhancing the overall quality and interpretability of the learned representations.
3
Background
The fundamental concept of the Denoising Diffusion Model is to gradually introduce small amounts of noise to clean imput sample S, transforming it into a representation, Sr, that can be sampled from a Gaussian distribution. This transformation follows an incremental progression: So + S1 + ... + Sr. The original data point So is drawn from the data distribution So ~ (()), while the conditional distribution q(s; S(-1) is assumed be Gaussian. Each noisy sample at time step t is generated using the formula:
S = B 11-BC6-1
(1)
where 3e represents the noise schedule, and ct-1 is sampled from a Gaussian distribution. By employing the notation a₁ = 1 Be and IT Oi, Se can be sampled in closed form
S. Vā + v1 - a.c.
(2)
Thus, the true posterior of the sample can be represented as q(Se-1 S) = -N( BrI), where für and å are computed using the following formulas:
V S+
Vo ōx-1) Sp, B, 1 at
1-a a1-1 Br.
Pe 1 az
(3)
- ae
In atlempting to reverse this process, approximate distribution is learned by pa | Si) (o B,I), where po represented as
1 (S.- B. Var VI-ai aq
p4
go (S,,))))
(4)
This representation allows the diffusion process to minimize the KT. divergence between the posterior and the prior distribution of samples (Ho et al. 2020). go(S,,t f) is a neural network that estimates the noise given the current state of the sample and the time step f. Therefore, sampling from po St)) allows us to directly estimate S as:
B ( go (S.,t) -_-30)0))))) VI á
1 Ss
+ 1- 1-a0-1 Bie.
a1-1
S
Vat-
1 ar
This process enables the gradual denoising of a sample that was originally drawn from noise distribution until reaches a clean data sample from the distribution of clean samples.
It is important to note that the term "time step" is used in two distinct contexts in this work: the definition of categorical time series and in the operation of diffusion models. In the context of a time series, "time step" refers to the specific time point at which each observation in the series collected. In contrast, within diffusion models, 'time step' denotes stage in the iterative process of adding noise to the imput embedding. The number of time steps in diffusion model controls the smoot of the diffusson process, determining how progressively an input is transformed into its corresponding noise representation.
4
The Proposed Model
Given time scries comprising r individual time steps, denoted as w = {wT w ,., WT where each time slep includes a vector w, E IIxx, our method learns to convert this series into a meaningful embedding representation. Each w, contains a set of q, categorical values. Our approach consists of two main components: data transformation module and a diffusion-basee representation learning module. The data transformation serves a learnable preprocessing step, performed independently before the initiation of the representation learning process.
4

---

Published in Transactions on Machine Learning Research (03/2025)
Forward path of DDPM
W X45'
Daca Transformation
'73''
'A67')
Denoising
S
'408'
'Y55'
'811'
Supervise
So
Zeem
Zeam while keeping it dense
Latent Space Navigation
Creating Spelling vector
d va
Learnable parameters
Figure 1: General pipeline of our method. The main objective is to extract meaningful representation, Zsems from a categorical time series The training process consists of two steps: first, transforming the categorical series into a continuous representation Sp = WAr, where is a multihor representation of w and A is a linear transformation; second, projecting it into a latent space, Esems and utilizing it to guide a conditional diffusion model. 7eem is supervised through Latent Space Navigation which allows this guidance to be dense enough while being supervised by sparse vector. The sparse vector provided by the Spelling information of the tokens in w which is constructed using the methodology illustrated in Figure 12 (Appendix).
4.1 Data transformation
The embedding layer is a fundamental component neural language modeling (Raffel et al. 2020 Liu et al. 2019 Devlin et al. 2018 Brown et al. 2020), functioning as a lookup table with learnable parameters. It associates each token within a text with continuous representation, effectively converting each token into fixed-size vector. However, this approach has limitations when dealing with categorical time series, where the number of lokens per time step varies and is not predefined. This variability can make it difficult to create consistent fixed-size vectors for different time steps. Furthermore, utilizing an embedding layer before a contimuons diffusion model requires the network to optimize for two distinct, interdependent tasks. In this scenario, the diffusion model receives an input that changes with each training iteration because the input of the diffusson model derived from the output of the embedding layer, which itself optimized during training, As a result, noise in the diffusion process is added to an input that varies throughout the training process. Consequently, the network may not simulate stable forward path during its denoising backward process, potentially leading to slow even lack of convergence.
To address this challenge, we propose an approach that learns to generate a fixed-size continuous representation for each step' of a time series, irrespective of the number of tokens in each step. Similar to the embedding layers, we begin by establishing a vocabulary comprising all unique tokens our dataset, assigning an index to each token. We then construct a binary vector for each time step of the series, where each element is set to one if there is token in that time step whose vocabulary index matches the index of that element. This method allows us to create multi-hot representation of the tokens present at each step of the time series. In contrast to traditional embedding layers, which typically position tokens close together based on their contextual co-occurrences, our approach emphasizes mapping one set of tokens close to another set when they differ by fewer than f tokens Each set represents all the tokens present. in a single time step of series. To facilitate this, we initialize a transformation function, A designed to map each multi-hor hot vector to a continuous code. One straightforward method to implement this transformation by using a matrix that performs a linear projection on the input sample.
8; = w; Atr
(6)
2In this work, empirically
5

---

Published in Transactions on Machine Learning Research (03/2025)
where w, is the multihot representation of w, x w, Multihot(w.). For simplicity, we omit the explicit mention of subject indices in the notation of wy. A subject index indicates the subject from which the i-th time step is selected. Air is the transformation matrix which is initialized by random values Afrit e RWx" W; E Lx 8 E R W is the number of unique tokens the vocabulary, and m is the size of the continuous embedding vector. We normalize w, to have a unit lz-norm before incorporating into this equation, aiming to mitigate the influence of the number of tokens on the performance of our algorithm.
To ensure that contextually similar sets of tokens mapped to nearby points, first need to construct pairs of token sets that contextually similar. For this purpose, propose to initially introduce set of minor alterations to Wi, by randomly changing some of its binary elements from ones to zeroes and vice versa. This approach ensures that the resulting set of tokens remains contextually similar to the original set, as most of the tokens are preserved. This resulting vector is denoted as w Subsequently, we aim to find a new Air that can map w close to &,, which has already been calculated as the embedding of w4- To achieve this, we seek an Ape" that satisfics the following equation.
y w weww Anew
s.t. |
1l2
(7)
The constraint ||Age||| = 1 is employed to normalize the Iransformation matrix, which prevents it from becoming too large or too small, thereby stabilizing the learning process. We define [ل This normalization facilitates the mapping of witw A and s; close to each other, while simultancously avoiding trivial solution for A that would transform all w; from different subjects to fixed 8;, irrespective of their actual values.
To express this in matrix form, we concatenate all w.ew vertically to create a matrix: Wenww [(w[ow)' (w (wnew)"" ..... (w (wRew)r]²" and similarly concatenate the different st values to form S
[sT sT ,...., sX] Then, Anow can be determined using the following equation:
(WWnew) Wneww + nT) -1 (WnewwTs. new
(8)
where IC is the number of different w, considered for training A,, This approach results in a new transformation matrix Anew that learns to position contextually similar sets of tokens close to each other in the embedding space.
To further the training process, we utilize Ape to map the w, values into the embodding space, thereby modify the w; values again and resolve equation
achieving the corresponding 8; values. Subsequently, (6) to obtain new AFew, repeating this process for specified mumber of epochs. Algorithm provides summary of this training process for our data transformation strategy.
The algorithm is designed to take a single multi-hot vector, Wi, as input. Instead of using real data to create this vector, we can generate efficiently by randomly placing ones in an all-zero vector. The number of ones corresponds to the average number of tokens that typically appear at cach time stop in our datasct. This approach allows us to sinulate all possible token combinations without the need for additional data collection, making it cost-effective way to cover a wide range of scenarios.

---

Published in Transactions on Machine Learning Research (03/2025)
Algorithm 1: Transforming categorical data to continuous for time series analysis with multiple tokens
per time step
Input: Multi-hot representation of tokens at each single time step, w I,,., K, Max epoch Output: The projection matrix Air which transfers each wi to continuous embedding vector 8;
1. Randomly initialize Atr
2. w/² |w Vi [L, W + W//
for e [0. Q do
3. Slight modification of Wt: Vi [1,K] W + Whew
4. w Tw 11
[1,
5. = W Atr
6. y: 1 [
7. Atr = ((Wnew) W + pl) (
8. Aer = A
end
4.2 Diffusion based representation learning
The main objective of this module is to project a time series, represented in continuous form as So E Brx", into an embedding space to obtain a compact representation, žeem = (So) E RT, To achieve this, we employ a conditional diffusion model trained in a self-supervised manner, where the condition and the target sample that the diffusion model learns to reconstruct are identical. However, the diffusion network is conditioned only on a compact representation of So, rather than on the entire time series. This encourages the network to capture as much relevant information as possible within this compact representation to enable accurate reconstruction of So. For simplicity in notation, denote the learned representation as 2eem which is then used as condition for the diffusion model.
In practice, we sample a noisy version of the time series So al a random time step / as follows: S. Va,So + -I- ate, We subsequently employ a denoising network, ge. to provide an estimation of the noise in the noisy sample of S.. With this value of noise, Io provides an estimation for the value of the noisy sample at time step 1: S
1 (S. 1 Vai
3t VI - - t, t,em)) +
1 0:-- B,C
1 ae
Therefore, minimizing the difference between the posterior and prior distributions directly leads to parameter optimization for Es, where Zsem (So). This enables the diffusion process extract meaningful information from Sp, as the learned features crucial for guiding the diffusion process in distinguishing between the noise and the core content of the noisy sample.
4.2.1 Cycle conditioning
In this section, we introduce the concept of cycle conditioned diffusion, which aims to encourage our diff'usion model to extract a multi-task representation from its conditioning samples. In traditional conditioning, diffusion models receive a clean version of a conditioning sample and use it to guide the denoising process toward the primary content of the conditioning sample. However, this method lacks a mechanism to ensure that the same latent represcntations are consistently extracted from contextually similar conditions that vary in fine details. As result, conventional conditioning does not guarantee that the extracted representations focus exclusively on the main content of the conditioning samples and not on extraneous details.
To address this limitation, we propose incorporating cycle consistency in the feature extraction process from conditioning samples. By enforcing consistency in feature extraction, this approach ensures that the diffusion model learns to produce a multi-task representation from the conditioning samples. This method not only guides the diffusion process using the fcatures of a clean condition but also applics the same guidance when the generated denoised sample is used as the condition for subsequent denoising steps. Consequently, we can predict additional estimate for the mean of the prior distribution of samples.
7

---

Published in Transactions on Machine Learning Research (03/2025)
Hg
cycle
1 (S: Va: v1
1
3.
go (S t, Ea 0z
1 vI- ai Vā.
S
vā ((1)))
(10)
where 6t-1 = 90 L, (oo())) By enforcing this cycle constraint, the model learns to capture the same content from both the clean and the noisy samples. This prevents the model from simply distorting the latent condition to match the dimensions S and then using S, add noise to this distorted sample. Instead, the model learns to perform the same guiding with different variations of the conditioning sample. Similar to vanilla DDPM, this new prior mean is drawn from minimizing KL divergence between the posterior distribtion of samples and the prior distribution using cycle consistent diffusion process.
4.2.2 Spelling supervision
One main disadvantage of cycle conditioning in diffusion models is the neglect of conditioning information. This occurs because the condition each sample changes with every training iteration. As a result, the network cannot. consistently use the conditioning information as guidance to reconstruct the sample. To prevent this issue, we propose supervising both the condition and the denoising process with additional information extracted from the tokens of the original time series before converting them to continuous form. This approach empowers the diffusion model to differentiale between the contributions of its conditioning information and that of the noisy sample, thereby preventing blind fusion of these inputs, which could lead to an inaccurate feature representation of the conditioning sample.
To do so, we begin by constructing a spelling vector for each time series, which quantifies the spelling patterns of its tokens. Initially, we assign a numoric value, greater than 9, each character in our database, i.e. assigning one number to each of the 26 letters in the English alphabet. Numeric values less than 10 are preserved to represent real numbers, enabling the encoding of tokens that contain both numerical and alphabetic characters. Each token is then represented as numerical code by sequentially concatenating the assigned values for its characters, starting from the left. Zero-padding is applied if a token has fewer characters than the maximum length of tokens in our database. Subsequently, we create zero-vector whose length matches the maximum possible numeric code in our dataset. Entries in this vector are set to one if there are tokens whose numeric code matches the corresponding index. This resulting vector referred to as the spelling vector. This process is graphically depicted in the right panel of Figure 1.
For this supervision, consider function that receives the latent condition 2sem Eg (So) and learns to predict the elements of the spelling vector. Let p Ma(zsem) denote the predicted values of the spelling vector. The objective of the supervision is to minimize the cross-entropy error between the elements of the spelling vector u and the predictions
Lipt Σ (p))
(11)
where uy is the r-th element of the spelling vector. To avoid the network from overlooking the condition, these spelling elements are additionally predicted when cycle conditioning the process, w, MM
1 S
ya.
A-aa Va. 1)-)
Leycle-epl epl =
Σ Uglog((zz)
(12)
4.2.3 Latent Space Navigation
While spelling supervision prevents denoising network from overlooking its condition, the effectiveness of this supervision is mostly hindered by the sparsity of information in the spelling vectors. Supervising the latent condition Zsem = (So) with a sparse vector yields a sparse latent condition, consequently reducing its effectiveness as a multi-task representation (Figure 2, right panel).
To avoid this issue, we propose to perform this supervision after navigating each latent condition in a set of meaningful directions (Figure 2, left panel). Our method involves learning set. of latent codes
8

---

Published in Transactions on Machine Learning Research (03/2025)
St
Denoising
St-1
St
Denoising
S
Supervised
C
C
Tarm
Supervised to indicate
Supervised
2am
Supervised to indicate
Figure 2: Direct supervision of the latent condition with spelling information from time series leads to a sparse latent space. This sparsity hinders its ability to serve as an effective vector for guiding the diffusion model or as compact, multi-task representation (right pancl). contrast, Latent Space Navigation individually guides each component of the latent space by navigating it in a specific direction. This approach allows for denser latent code, while supervised by the sparse spelling information from the time series (left panel).
dxp Rmm F {1, v each representing an initial direction towards one of the elements in the spelling vector, Ug. Consequently, Zaem (S(o) is mapped to an individual target point in the latent space, Asem- before being employed to predict the x-th element of the spelling vector, Uz!
Ssem-hu,
2sem + dzyg
(13)
This approach allows the network to avoid directly learning both the direction and magnitude of all Ue values on the latent condition (So) Instead, the network Icarns the magnitude of cach uz after traversing in the correspouding direction. This strategy facilitates the supervision of =sem = (a(Sg) using a sparse spelling vector while ensuring sufficient density function as an effective multi-task representation. The same navigational strategy is employed prior to predicting each spelling element Uy from the cycle-conditioned model:
VI- - at Et-1)
Zeaem-Hu. Eg( S.
das
(14)
Va
where f-11 = go Eo(Sa)). Considering this navigation, the supervision terms Equations (11) and (12) are reformulated as follows:
Lapt Σ (Z)eemm-uu)))
(15)
Σ (Mlog(MM(essemmuu..
Leyole-spy oopl
Detailed configurations of our method including the configuration of Eg and Mg are provided in Appendix
4.2.4 Learning
We train our model in a self-supervised manner, where the condition and the input to the forward path the diffusion model are identical. The overall loss function minimizes the MSE loss between the predicted and the true noise, considering both the base and the cycle conditioning of the denoising network. Additionally, it includes the prediction loss for spelling errors associated with these two conditioning mechanisms.
In practice, we leverage Multi-Step Optimization approach, which allows the introduction of extra loss functions alongside the MSE loss of the diffusion models with minimal impact on their performance, Specifically, we alternate between optimizing the MSE and cross-entropy losses. The cross-entropy losses, Lepi and Leycle-sph are only applied any of the MSE losses exceed a threshold of 0.05. During the multi-step optimization, we
3p n where the number of possible characters (letters numbers) the dataset, and the length of the longest. token

---

Published in Transactions on Machine Learning Research (03/2025)
perform single forward pass to compute both the DDPM's MSE losses and the cross-cntropy losses. In the first step, we update the model parameters with respect to the MSE losses while freezing the components Ma and dg affected by the cross-entropy losses. In the next step, we update the model with respect to the cross-entropy losses, freezing the components relovant to the MSE losses.
4.2.5 Inference
Inference strategies may vary depending on the downstream task. When a substantial amount of training data with corresponding labels is available, we can directly extract representations from the latent space of the conditioning encoder, Eg (So). These representations are sufficient for effectively training a downstream classifier.
Conversely, in scenarios with limited number of training samples, it is advantageous to leverage the generative capabilities of our method. In this case, representations are derived from Eg(S,, f, Ee( Va. S, where S, va,So + vI- a,,, and is randomly sampled from Gaussian distribution.
yI-ãe e--1)),
vā,
5 Experiments
Our experiments are conducted using two well-known databases, MIMIC III (Johnson et al. 2016 and MIMIC IV (Johnson et al. 2023). each containing records of medical activities from various patient visits to a hospital. We construct. a time series for each patient in these databases. These series are built. from the diagnosis codes of the patients, as they have minimal missing values compared to other activity records in these databases. The diagnosis codes collected during each patient's visit, form feature vector, representing one step of time series. Each feature vector is truncated to length of 13, enhancing the compulational efficiency of our analysis. We only consider the first four visits of each patient. While MIMIC IV contains both ICD-9 and ICD-10 codes, MIMIC III only includes ICD-9 codes. To ensure consistency, we truncate the codes to their first three characters and prefix them with either '9' or '10' to denote their coding system. This ensures the maintenance of relevant information while excluding unnecessary details from our analysis.
Data split: Most of our experiments are conducted on MIMIC IV, where we randomly select 126,736 time series for training our representation learning method. The evaluation sets for different tasks are composed of another 31,684 randomly selected samples, which are used to evaluate the performance of the learned representations various downstream tasks. For each task, positive cases include patients diagnosed with specific disease, while negative samples are randomly selected from the remaining patients, ensuring an equal number to the positive cases. Therefore, not all samples are used for training the downstream classifiers. For MIMIC III, all samples serve as training data for the represontation learning model. The training and the evaluation of the reanstream task are then conducted using the same split of the MIMIC IV.
Preparing data for the downstream task: After training our diffusion model, the training and test samples are separately projected into the embedding layer of this representation learning framework. The representations of the training samples are then utilized to train downstream classifier. This classifier learns to predict a specific disease within a time window of 4 months before it is diagnosed by physicians in the real world. Patients who are diagnosed with the disease before this specified time window are excluded from the training process of the downstream classifier. For the downstream classifier, we employed Support Vector Machines (Cortes & Vapnik 1995). Gaussian Process (Williams & Rasmussen) 2006), Random Forest. (Breiman 2001), Adaboost (Freund et al. 1999), and Gaussian Naive Bayes (Gaussian which are selectod based on their proven performance in similar lasks.
Evaluation metrics For evaluation, we the following metries: AUROC: This measures the ability of a classifier to distinguish between positive and negative classes across different threshold values. AUPRC: This represents the trade-off belween Precision and Recall of a classifier, indicating its ability to correctly classify positive samples as positive. Cohen's Kappa: This measures how often classifiers correctly classify positive cases, while also considcring that some correct predictions might occur by chance. Macro-average Fl-score: This is an average F1 score for each of the positive and negative classes, treating them equally regardless of the number of positive and negative cases among the evaluation samples.
10

---

Published in Transactions on Machine Learning Research (03/2025)
Heart failure prediction
Our method
SwatelTab
Recon Tab
SCARE
MLM
Cut M.M.--MMMM MLM
Em--AuCu Con
Wa et al
2021
Chen et al. 2023
Bahri et al. 2021 Devlin et al. 2015
Devlin ct a. 2018
Chen ot a.
2024b
AUROC AUCPR,
ALROC
ATCPIt
ALROC AUCPP AUROC ACCPI
AUROCT
AUCPR
AUTOC
AUCPR
AUROC
ACCPI
Support Vector
91.8..6:
12.49.as 90.39.0.75
a9. 61.79u0.10 02.54
...39 06.48a77
09.49.0.4 03.15
74.34.0.52
04 44.46.88
77.00.0.4
05.11.0%
Gauesian Proces 00.96-042 13.70-ss 90.23-0..1 10.09-48, 87.57-045 08.86-4%
79.85-9.16
19.13-69
64.10 10.0055.55 02 86.6 76.63.0.56
05.01-4%
76:22-049
07 07-69-15
Random Forest 89.48.e.r
14.21.95
86. 12.n.s 07. 47.4
84.0
06.72
a0.87
06.46aa5
07.63.0.4
03.31
68.09ue.2
03. 03.48.0%
77.33.074
05.26a38
Ada Boost
85.12-073
68.62 80.53=0:: 06.51-6% 65.87 03.01:4%
78.04=0.1 05.56a7%
64.25-0.47 02.58
66.62am45
03 17.t 71.21=0.00
03.87at%
Gaussion NB
89. 80.20-0.2 08.4 -7% 16.16-06 (6.28--2 5% 64. 18-0:55 02.79 6% 77.93-965 05.46-5% 60.58-0.8 02.544055 70.04-021
03 63--63-5 74.95-083
04.67-15
Av Kappa BP1 Av Kappa Av BFT Av Kappa Av. BFI 20 Kappa /o [F1 Av Kapon BFT Av Kappa AF.
Kappa
AF.
18.84-04652.2 26.4 15.0 42.0o: 12.72.1 45.15.0 16.0440.7 49.40. 14.222..55 45, 18.0.5 15.74.0.51 48.31.085
16.09-033
48.92.074
Langg cancer prediction
Support Vector 90.51-0st 0.54.M
55.98-0.55
0.42-4%
00.25-0.2
0.15.a% 86.75 66.75-0.55 0.46-115
83.47-0.7
0.34-9%
8261-0a
0.42-4%
50.44-05s
0.30-s%
Goussian Proces 88.79.036
0.5G
80.26.0.7
0.43
63.71.0.72
0,18 85, 18.061 0.38.6x
80.14.0.77
0.31448
82.94.077
0.37,2%
80.01.039
0.31.48
Random Forest 91.35.0.45
0.50am
85.91a043
0,36amx
67,39
0,31
82.82=0.3 0,37.48
81.19
0.38
80. 60.14
0.33a:%
82.19.0:3
0.35G125
Ada Boost
86.60-057
0.49.%% 82. 12-12...
0.31
50.49.075
0.13ea% 85.57-9.7 0.36-7%
81,25-0.45
0:24-4%
8181--022
0.31-3%
77.19-000
0.22-78
GoussianNB
84.39-077
0.43.mx 88.32.0.15
0.4448% 62.35-045
0.14.88 84.33-9.67
0.33-58 82.12-0.16
0:29.48
80. 60.14.075 0.24.6%
81,17-032
0.36-48
Av Kappa RP1
Ar. RF1
Av: Kappa Rr.
Av Kappa Br.
Av Kapgo
Ax BF1 Av: Kappa
Av Ko.Kopp AF.
0.94.08 45.87aam 0.85n08 43.33.00r 0.33 41.59aam 0.70,7% 49.41.0ss 0.23.18 40.06.0.2
0.77a55 43.44.047
0.64a008 44.19.08
Table 1: Performance comparison between our method and the state-of-Che-art on heart failure and lung cancer prediction
5.1 Performance Evaluation in Healthcare Applications
In this section, we evaluate the performance of our diffusion-based method in learning representations that are effective in various downstream tasks. Specifically, tasks that are completely unrelated and need highly expressive representations for accurate identification. Therefore, demonstrating the capability of a single representation in predicting both of these tasks can prove the multi-task utility of that representation. Our evaluation includes comparisons with the state-of-the-art methods in representation learning from categorical data: SwitchTab Wu et al. 2024 ReCon'Tab (Chen) et al. 2023) and SCARF (Bahri et al. 2021).
consider heart failure and lung cancer as two distinct
Additionally, we compare our results with the representations learned by the embedding layer and one layer before the last of the BERT model (Devlin et al. 2018), pretrained using Mask Modeling. These methods are respectively denoted as En-MLM and OutBERT-MiN
Furthermore, we compare our method with the same BERT structure, trained using the self-supervised learning strategy proposed by Chen et al. (Chen et al., 2024b). referred to as Em-AuCon. For each method, we consider a four-month window preceding discase diagnosis, providing an optimal limeframe for initialing medication upon correct disease prediction, thus reducing the risk of fatality associated with these diseases. The representation dimension of our method, along with the models that utilize BERT as their backbone architecture, are set to 400. For SwitchTab(W et al. 2024). ReConTab (Chen et al. 2023). and SCARF (Bahri et al. 2021 follow the same protocol described in their original papers, resulting in representation dimensions of 52, 256, and 256, respectively.
As can be observed, our method outperforms all other strategies across all the metrics. For instance, in heart failure prediction, it surpasses SwitchTab's best results by 1.42%, which itself is the leading performance among the competing strategies. Even the lowest performance of our method, achieved by Adaboost, is 94.13% of the best results achieved by other competitors, demonstrating consistent performance of our method despite using different classifiers for the downstream task. In the case of the lung cancer prediction, our best. result outperforms the best results of the other competitors by 3.3%, while the lowest performance of our method still achieves 95.57% of this best result.
5.2 Probabilistic representation learning
In this section, we demonstrate the effectiveness of our method in generating different meaningful represen- tations for a single input sample. This capability is particularly valuable in applications where there is a limited number of labeled samples to train a downstream classifier.
In this case, instead of directly obtaining representations from (So),, we acquire the representations from VI-B
Eg
LS S.
Et-1), where 6 is the estimated noise of the noisy sample St, estimated by ge).). For V
11

---

Published in Transactions on Machine Learning Research (03/2025)
this, we consider B random St with 3e value between 1e and 0.5e for each positive patient to generate various representations of the disease. Notably, only one forward pass is needed to compute each 6t-1 from S, with no iterative process involved, making our method exceptionally fast in generating novel representations of samples. By employing this simple strategy, we ensure there is enough number of training samples for the
downstream classifiers, which is B 15 or 30 times more training samples compared to our baseline model.
Our method (Prob. #16)
Our method (Prob. #30)
Out method (Prob. #30)
Our method (Prob. #60)
AUROC
AUCPR
AUROC
AUCPR
AUROC
AUCPR
AUROC
AUCPR
Support Vector 92 22.38t0.58 6a
1632290%
83 80 58
19.1445%
Support Vector 93. 13.17:0.43
0.74=4%
83.75400
0.8
Gaussian Process 92.03+00
17.00++4%%
#2 42
17
Gauseian Process 90.82---557
0.53-1%
90.944+0335
0.56+-%
Random Forest 918440.34
16.1420%
92 33.2..65
16. 16.32455
Random Forest 92.68=0.72
0.70=7%
02.8540.22
0.714
Boud GaussianNB 54
Laca
11.5316%
Lo..as
11.96.19%
Boos 0..9
0.51..4%
90.35.1.47
L.5
12.75+6%%
90. 90.32++0.59
14.17+
GanssianNB
89. 89.8--0511
0.50-55
50,10+606
0.61 +0%
Table 2: Probabilistic representation learning using our Table 3: Probabilistic representation learning using our method for heart failure prediction. The number of syn- method for lung cancer prediction. The number of synthe- thesized samples for each real patient is 15 and 30, respec- sized samples for each real patient is 15 and 30, respec- tively.
tively.
Results: As demonstrated in Tables 2 and 3, our strategy clearly outperforms the basclinc model in both the downstream lasks. For instance, we observe an overall increase in AUROC for both the lasks. Specifically, the performance of our method improves from 91.81% to 92.38% for heart failure prediction and from 91.35% to 93.17% for lung cancer prediction, with just 15 additional representations for each positive sample. In contrast, achieving AUROC scores of 93.80% and 93.73% for these tasks requires 30 times more training samples for the downstream classifier.
5.3 Our data transformation strategy Embedding layers
In this section, evaluate the ability of Algorithm 1 to capture meaningful information from sequence of categorical tokens, comparing its performance to that of the BERT embedding layer. Following strategies from prior works (Kennard et al., 2016 Bojanowski et al. 2017 Pennington et al. 2014), which propose to evaluate the embedding layers based on their capacity to preserve word similarities in the embedding space, we examine how effectively each embedding layer maintains sentence-level similarity within its embedding layer. This property is especially relevant when outputs of the embedding layers are further used as inputs to diffusion models, where each sample is treated as a distinct point in a high-dimensional space rather than as composite of multiple subspaces.
To evaluate this capability, we first concatenate the textual descriptions of all ICD codes associated with each visit in every time series within our database. We then randomly select two visits from different series and use a pre-trained large language model to obtain embeddings for these concatenated descriptions. By calculating a cosine similarity score between these embeddings, we mcasure the ability of each embedding layer to capture the semantic similarity between two distinct sets of ICD codes. For evaluation, we transform each set of ICD codes (corresponding to individual patient visits) into a continuous representation using both our method and the BERT embodding layer. We then calculate cosine similarity score for cach embedding, which enables us to assess how effective each method correlates with the similarity scores derived from the ICD code descriptions. We adopt two distinct strategies for learning parameters in BERT embedding layer. The first approach is through Masked Language Modeling (MLM), a conventional method commonly utilized for self-supervised training of the embedding layers in large language models.
Additionally, we explore an extension of the method proposed by (Chen et al., 2024b). which is a multi-task strategy for self-supervised representation learning from images. However, we adapted this strategy for categorical data to enable the capture of relationships between different tokens in the immediate output of the embedding layer. In the case of masked modeling, we extract the embedding data from two different layers: one from the conventional embedding layer and another from the layer preceding the softmax operation at the end of the network. This is because masked modeling does not inherently encourage the embedding layer to accumulate the contextual relationships between different words in each embedded token. Instead, it captures these relationships at the end of the network, after the transformers of BERT, where it needs to estimate a masked token from a set of surrounding tokens, and therefore needs to accumulate this surroundigg information for each token.
12

---

Published in Transactions on Machine Learning Research (03/2025)
Table 4 presents comparison between these four strategies using Spearman correlation. Our method shows high positive correlation across different embedding dimensions, suggesting a strong association between its similarity scores and the ground truth similarities. In contrast, the BERT embedding layer trained by MLM returns the lowest correlation, aligning with our premise that masked modeling does not inherently capture the contextual relationship between different tokens.
Dimensionality
40
100
200
300
400
700
1000
1500
Em-MLM
0.14+
0.17++0.06
0.183
08
0.15+
0.17+#0.06
0.19440.04
0.20++0.08
Out-MLM
0.33-10.08
0.34:10.08
0.34:10,008
0.36:10.10
0.36:10.08
0.37:10.05
0.38:10.07
0.39:10.00
Em--onAu
0.18
0.21
0.22::0112
0.23
0.24
0.27
0.28
0.29:
0.38:t0.12
0.39
0.39
0.41+t0.08
0.42:t0.09
0.44+t0.08
0.46-t0.09
0.45+t0.05
Algorithm
0.42++0.11
0.43-+0.07
0.44++0.09
0.45-+-0.06
0.45-+0.07
0.49++0.08
0.51 +
0.51
Table 4: Performance comparison between our method and the BERT embedding layer on the MIMIC-IV dataset based on Spearman correlation over 20 runs.
5.4 Ablation study
We also conducted an ablation analysis focused on answering the following key questions:
Does cycle conditioning aid in extracting more versatile representations? To investigate this, trained two different. versions of our method: one incorporating cycle conditioning loss and the other without it, while preserving the spelling loss functions in both versions. This design allows us to evaluate the specific contribution of the cycle conditioning loss to the overall performance of our method. As shown in Table 5, the network without cycle conditioning achieves significantly lower accuracy in the heart failure prediction task, demonstrating the effectiveness of this loss function in enhancing the performance of our method.
2. Is spelling supervision necessary? The conditioning loss function may encourage the network lo extract the same representation from both clear conditioning sample and its noisy version, suggesting it could be sufficient for effective representation learning. To explore this, we trained our model using the cycle conditioning loss but without the spelling loss functions. To prevent the network from extracting identical representations from all conditioning samples, we normalized this layer to have the same standard deviation as the input conditioning sample. The results of this comparison are also presented in Table As observed, without the spelling loss function, our model struggles to capture meaningful representations from the conditioning sample, which hinders its ability to effectively estimate heart failure as the downstream task.
Full model
w/o eycle conditioning
w/o spelling supervision
AUROC
AUCPR
AUROC
AUCPR
AUROC
AUCPR
Support Vector
91.81=0061
12.49.33%
89.47
12.23178
63.21
02.78:99
Gaussian Process
90.96-00662
13.70++59
87.93+00.72
099022+-9%
62.16
02.14+-9%
Random Forest
89.48=0.37
14.214:%%
88.51
12.61
64.211:047
02.8413%
Ada Boost
85. 12-3773
08. 62.16%
83.77 83.77:10.67
07. 17.
58. 28.12:bo.4 43
02. 02.11:111%
GBussianNB
89.20:10.76
08. 08.41.77%
87 87.23:10.58
08.09-14%
58.69:10.84
02.26.18%
Table 5: Ablation study proposed method where the effectiveness of its two main components, cycle conditioning and spelling supervision are evaluated
6 Conclusion
In conclusion, we introduced a novel diffusion-basee method for representation learning in categorical data. We proposed that cycle conditioning of diffusion models can effectively extract meaningful and versatile
13

---

Published in Transactions on Machine Learning Research (03/2025)
information from the conditioning samples. Additionally, we introduced a novel strategy to supervise the condit ioning process with extra information about the spelling of tokens, thus preventing the overlooking of the conditioning samples in the denoising process. To enhance the robustness of the representations, we proposed supervising the latent condition through linear navigation in its latent space. We conducted a comprehensive set of experiments to demonstrate the effectiveness of our diffusion model in extracting meaningful representations from categorical data. However, as with any approach, this method presents certain limitations, including significant computational demands inherent to diffusion-based models. Future research may focus improving computational efficiency, potentially by integrating faster generative modeling techniques for representation learning tasks. Furthermore, evaluating the model's performance on wider range of databases would offer valuable insights into its robustness and adaptability. Specifically, exploring its sensitivity to diverse categorical datasets, beyönd the medical domain, could provide a deeper understanding of its capability to learn representations in different contexts and inform its application to medical time series representation learning.
Acknowledgements
This study was part of the AIR Lund (Artificially Intelligent Use of Registers at Lund University) research environment and received funding from the Swedish Research Council (VR; grant no. 2019-00198). It also received funding from the Knowledge Foundation (CAISR Health with grant 20200208 01H). The funders had role in study design, data collection and analysis, decision to publish, preparation of the manuscript.
References
Rameen Abdal, Yipeng Qin, and Peter Wonka. Image2stylegan: How to embed images into the stylegan latent space? In Proceedings of the IEEE/CVE international conference on computer vision, pp. 4432 4441, 2019.
Arash Afkanpour, Vahid Reza Khazaie, Sana Ayromlou, and Fereshteh Forghani. Can generative models improve self-supervised representation learning? ar preprint arXiv:2403. 059666 2024.
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34:17981-1777 17993,, 2021.
Dara Bahri, Heinrich Jiang, Yi Tay, and Donald Metzler. Scarf: Self-supervised contrastive learning using random feature corruption. Xriie preprint arXiv:2106.1511 2021.
Ankan Kumar Bhunia, Salman Khan, Hisham Cholakkal, Rao Muhammad Anwer, Jorma Laaksonen, Mubarak Shah, and Fahad Shahbaz Khan. Person image synthesis via denoising diffusion model. In Proceedings of the IEEE/CV Conference on Computer Vision and Pattern Recognition, pp. 5968-5976, 2023.
Piotr Bojanowski, Edouard Grave, Armand Joulin, and Tomas Mikolov. Enriching word vectors with subword information. Transactions of the association for compulational linguistics, 5:135 146, 2017.
Lco Breiman. Random forests. Machine learning, 45:5-32, 2001.
Tom Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 13:1877-1901, 2020.
Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. Advances in neural information processing systemss, 33:9912 9924, 2020.
Pei Chen, Soumajyoti Sarkar, Loonard Lausen, Balasubramaniam Srinivasan, Shong Zha, Ruihong Huang, and George Karypis. Hytrel: Hypergraph-enhance tabular data representation learning. Advances in Neural Information Processing Systems, 36, 2024a.
14

---

Published in Transactions on Machine Learning Research (03/2025)
Suiyao Chen, Jing Wu, Naira Hovakimyan, and Handong Yao. Recontab: Regularized contrastive representa- tion learning for tabular data. arXiv preprinl ar 18i:22300188411, 2023.
Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, and Pieter Abbeel. Infogan: Interpretable representation learning by information maximizing generative adversarial nets. Advances in neural information processing systems, 29, 2016.
Xiaokang Chen, Mingyu Ding, Xiaodi Wang, Ying Xin, Shentong Mo, Yunhao Wang, Shumin Han, Ping Luo, Gang Zeng, and Jingdong Wang. Context autoencoder for self-supervised representation learning. International Journal of Computer Vision, 132(1):208 223, 2024b.
Corinna Cortes and Vladimir Vapnik. Support-vector networks. Machine learning, 20:273 297, 1995.
Jacob Devlin, Ming-Wo Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. orXiv preprinl arXiv:1810.04880 2018.
Sander Dieleman, Laurent Sartran, Arman Roshannai, Nikolay Savinov, Yaroslav Ganin, Pierre H. Richemond, Arnaud Doucet, Robin Strudel, Chris Dyer, Conor Durkan, et al. Continuous diffusion for categorical data. Xiviv preprint arXic:2211.15000 2022.
Tianyu Du, Luca Melis, and Ting Wang. Remasker: Imputing tabular data with masked autoencoding. aiv preprint ar (it::2309.13790 2023.
Yoav Freund, Robert Schapire, and Naoki Abe. A short introduction to boosting. Journal-Japanese Society For Artificial Intelligence, 14(771-780):1612, 1999.
Zhujin Gao, Junliang Guo, Xu Tan, Yongxin Zhu, Fang Zhang, Jiang Bian, and Linli Xu. Difformer: Empowering diffusion models on the embedding space for text generation. arXiv preprint arXiw:2212.09412 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in newral information processing systems, 33:6840 6851, 2020.
Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling, Argmax flows and multinomial diffusion: Learning categorical distributions. Advances in Neural Information Processing Systems, 34:1 24:12454-12488 2021.
Drew A Hudson, Daniel Zoran, Mateusz Malinowski, Andrew K Lampinen, Andrew Jaegle, James L McClelland, Loic Matthey, Felix Hill, and Alexander Lerchner. Soda: Bottleneck diffusion models for representation learning. In Proceedings of the IEEE/CVF Conference on Compuler Vision and Pallern Recognition, pp. 23115-23127, 2024.
Naoto Tnoue, Kotaro Kikuchi, Edgar Simo-Serra, Mayu Otani, and Kota Yamaguchi. Layoutdm: Discrete diffusion model for controllable layout generation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pollern Recognition, pp. 10167 10176, 2023.
Gautier Izacard, Mathilde Caron, Lucas Hosseini, Sebastian Riedel, Piotr Bojanowski, Armand Joulin, and Edouard Grave. Towards unsupervised dense information retrieval with contrastive learning. arXiv preprint arXiv:2112.09116 2(3), 2021.
Taeuk Jang and Xiaoqian Wang. Difficulty-bsedd sampling for debiased contrastive representation learning. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 24039 24048, 2023.
Hong Jia, Yiu-ming Cheung, and Jiming Liu. new distance metric for unsupervised learning of categorical data. IEEE transactions neural networks and learning systems, 1766):106,- 2015.
Songlei Jian, Guansong Pang, Longbing Cao, Kai Lu, and Hang Gao. Cure: Flexible categorical data representation by hierarchical coupling learning. IEEE Transactions on Knowledge and Data Engineering, 31 (5):8533 866, 2018.
15

---

Published in Transactions on Machine Learning Research (03/2025)
Alistair EW Johnson, Tom J Pollard, Lu Shen, Li-wei H Lehman, Mengling Feng, Mohammad Ghassemi, Benjamin Moody, Peter Szolovits, Leo Anthony Celi, and Roger G Mark. Mimic-iii, freely accessible critical care database. Scientific data, 3(1):1 9, 2016.
Alistair EW Johnson, Lucas Bulgarelli, Lu Shen, Alvin Gayles, Ayad Shammout, Steven Horng, Tom J Pollard, Sicheng Hao, Benjamin Moody, Brian Gow, et al. Mimic-iv, freely accessible electronic health record dataset. Scientific data, 10(1):1, 2023.
Neha Nayak Kennard, Gabor Augeli, and Christopher D Manning. Evaluating word embeddings using representative suite of practical tasks. In Proceedings of the 1st workshop evaluating vector-space representations for nlp, pp. 19-23, 2016.
Phuc H Le-Khac, Graham Healy, and Alan F Smeaton. Contrastive representation learning: A framework and review. Teee Access, 8:193907-193934, 2020.
Qiude Li, Shengfen Ji, Sigui Hu, Yang Yu, Sen Chen, Qingyu Xiong, and Zhu Zeng. A multi-view deep metric learning approach for categorical representation on mixed data. Knowledge- Based Systems, 260:110161, 2023.
Based
Jay Chichen Liao and Cheng- To Li. Tabgsl: Graph structure learning for tabular data prediction. arxiv preprinl 23v: 15843, 2023.
Zhenghao Lin, Yeyun Gong, Yelong Shen, Tong Wu, Zhihao Fan, Chen Lin, Nan Duan, and Weizhu Chen. Text generation with diffusion language models: A pre-training approach with continuous paragraph denoise. In International Conference on Machine Learning, pp. 21051 21064. PMLR, 2023.
Yinhan Liu, Myle Out, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zeettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprinl 19o: 11692, 2019.
Kushal Majmundar, Sachin Goyal, Praneeth Netrapalli, and Prateek Jain. Met: Masked encoding for tabular data. Xiv preprind arXin:2206.0888 2022.
Sarthak Mittal, Korbinian Abstreiter, Stefan Bauer, Bernhard Schölkopf, and Arash Mehrjou. Diffusion based representation learning. In International Conference Machine Learning, pp. 24963 24982. PMLR, 2023.
Jeffrey Pennington, Richard Socher, and Christopher D Manning. Glove: Global vectors for word representa- tion. In Proceedings of the 2014 conference on empirical methods in natural language processing (EMNLP), pp. 1532 1543, 2014.
Rudra PK Poudel, Harit Pandya, Stephan Liwicki, and Roberto Cipolla. Recore: Regularized contrastive representation learning of world model. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 22904-22913, 2024.
Konpat Preechakul, Nattanat Chatthee, Suttisak Wizadwongsa, and Supasorn Suwajanakorn. Diffusion autoencoders: Toward a meaningful and decodable representation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pottern Recognition, pp. 10619-10638 2022.
Alec Radford, Luke Metz, and Soumith Chintala. Unsupervised representation learning with deep convolutional generative adversarial networks. arXiv preprint arXi::1511.0644, 2015.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, (1400) 67, 2020.
Machel Reid, Vincent J Hellendoorn, and Graham Neubig. Diffuser: Discrete diffusion via edit-based reconstruction. arXie preprinl ar Xrxi:22200. 11.:2220.166866 2022.
Nikolay Savinov, Junyoung Chung, Mikolaj Binkowski, Erich Elsen, and Aaron van den Oord. Step-unrolled denoising autoencoders for text generation. arXiv preprint or Xrx2::1 12. 2.0644, 2021.
16

---

Published in Transactions on Machine Learning Research (03/2025)
Talip Ucar, Ehsan Hajiramezanali, and Lindsay Edwards. Subtab: Subsetting features of tabular data for self- supervised representation learning. Advances in Newral Information Processing Systems, 34:18853 1883-118655 2021.
Siyu Wang, Xiaocong Chen, Quan Sheng, Yihong Zhang, and Lina Yao. Causal disentangled variational auto-encoder for preference understanding in recommendation. In Proceedings of the 46th International ACM SIGIR Conference on Reseorch and Development in Information Retrieval, pp. 1874-1878, 2023a.
Yingheng Wang, Yair Schiff, Aaron Gokaslan, Weishen Pan, Fei Wang, Christopher De Sa, and Volodymyr Kuleshov. Infodiffusion: Representation learning using information maxinizing diffusion models. In International Conference on Machine Learning, pp. 36336-36354. PMLR, 2023b.
Christopher KI Williams and Carl Edward Rasmussen. Gaussian processes for machine learning, volume 2. MIT press Cambridge, MA, 2006.
Chuhan Wu, Fangzhao Wu, Tao Qi, Yongfeng Huang, and Xing Xic. Fastformer: Additive attention can be all you need. Xiv preprint or orXiv:2100 10.. 20:00408.0904, 2021.
Jing Wu, Suiyao Chen, Qi Zhao, Renat Sergazinov, Chen Li, Shengjic Liu, Chongchao Zhao, Tianpei Xic, Hanqing Guo, Cheng Ji, et al. Switchtab: Switched autoencoders are effective tabular learners. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp.
2024.
Johnathan Xie, Yoonho Lee, Annie S Chen, and Chelsea Finn. Self-guided masked autoencoders for comain-agnostic superviedd learning. arXiv preprint arxi::2002. 14799, 2024.
Jiahuan Yan, Bo Zheng, Hongxia Xu, Yiheng Zhu, Danny Chen, Jimeng Sun, Jian Wu, and Jintai Chen. Making pre-trained language models great tabular prediction. arXiv preprint 2r8i::4433.08841, 2024.
Run Yang, Yuling Yang, Fan Zhou, and Qiang Sun. Directional diffusion models for graph representation learning. Advances in Neural Information Processing Systems, 36, 2024.
Xingyi Yang and Xinchao Wang. Diffusion model as representation learner. In Proceedings of the IEEE/CE /VF International Conference on Computer Vision, pp. 18938 18949, 2023.
Jiaxin Yin, Yuanyuan Qiao, Zitang Zhou, Xiangchao Wang, and Jie Yang. Mcm: Masked cell modeling for anomaly detection in tabular data. In The Twelfth International Conference on Learning Representations, 2023.
Jiaxuan You, Xiaobai Ma, Yi Ding, Mykel J Kochenderfer, and Jure Leskovee. Handling missing data with graph representation learning. Advances in Neural Information Processing Systems, 33:19075 19087, 2020.
Zhongqi Yue, Jiankun Wang, Qianru Sun, Lei Ji, Eric I Chang, Hanwang Zhang, et al. Exploring diffusion time-steps for unsupervised representation learning. arXie preprint arXiv:2401.11430, 2024.
Mingtian Zhang, Tim Z Xiao, Brooks Paige, and David Barber. Improving vae-based representation learning. arXiv preprint ar 2rx::2205.145333 2022a.
Zijian Zhang, Zhou Zhao, and Zhijie Lin. Unsupervised representation learning from pre-trained diffusion probabilistic models. Advances in Newral Information Processing Systems, 35:22117-22130,, 2022b.
Chengzhang Zhu, Longbing Cao, Qiang Liu, Jianping Yin, and Vipin Kumar. Heterogeneous metric learning categorical data with hierarchical couplings. IEEE Transactions on Knowledge and Data Engineering, 30(7):1254 1267, 2018.
17

---

Published in Transactions on Machine Learning Research (03/2025)
A Implementation details of the BERT embedding layers
In our study, utilize BERT the foundational architecture to evaluate the performance of embedding layers. BERT's embedding layer is initially designed for neural language modeling, which typically requires large token vocabulary. To adapt it for our medical application, which has smaller vocabulary, we introduce a fully connected layer before the embedding layer. This modification allows us to control the size of the embedding space and identify the most suitable dimensions for our application. Additionally, we add linear layer at the end of the network generate logits vector for each index in the output sequence. The length of this logits vector corresponds the total number of tokens our medical vocabulary. We follow the standard masked modeling procedure by randomly masking 15% of the ICD codes, where 80% of the masked codes are replaced with predefined token, 10% are substituted with other ICD codes, and 10
We use the same training dataset as for our representation learning model. The learning rate is set to 1c-5, we employ the AdamW optimizer, and use cross-entropy loss to assess the estimated indices for the masked tokens.
255-3256 4:3
J55 356-)256,48 8:
556-9256, 8:3
51000000
256-9258,
258-2256
200-4 K-J
( Zeom)
sem
FastForner
Cone
Batch Norm
RelU
Conv 3D
Bartch
Retu
Cony
Batch Norm
Ri
Combination
Module
Self altf QUKVV
Combination
Module
Conv
RetU
Cony 10
Batch Norm
ReeUU
Corv
Batch Norm
Red
Conv 10
S
Latent Space Navigation
dig.
Com 1D
ReeUU
Com 1D
ReilU
con 1D
ReeU
Conv 1D
Ma (Zeem dia)
Pastoomer
1000--1000 FartFormer
Cony 1D
Fest
34, used
Conv
ReeU
Zyem
225-92.-7
suspervise Ma
1500 9253
4-32% K:
236-9256 K
2.50
Combination Module
Positional encoding
Fully commecti nonecctd
Group norm
Group norm
1D
ppp()
SUU
SLU
Conv
Figure 3: Implementation details of our method: For simplicity, we avoid illustrating the cycle conditioning part, which necessitates reusing St-1 as the conditioning sample instead of C.
B Our method vs other diffusion based strategies
We conduct an additional experiment evaluate the effectiveness of our method in comparison to other strategies that utilize diffusion modeling for representation learning. Specifically, compare approach
with those presented in (Mittal et al. 2023 Preechakul et al. 2022 Wang et al. 2023b). It is important to note that we do not include comparisons with Zhang et al. 2022b Yang & Wang 2023), as their methods focus on extracting representations from a pre-trained diffusion model without using a conditioning mechanism.
To ensure a fair comparison, we use the same structure for the conditioning encoder and denoising components in all the diffusion models. The only variations are in the transformation of categorical features to continuous ones and in the formulation of loss functions, where different supervision methods are
Fair comparison:
18

---

Published in Transactions on Machine Learning Research (03/2025)
256-2266 K:
256-9256 K:
256-9256, K2.
256- K:3
100-11009
256-2255
256-H4 K:3
(x Zgom)
Zsem
Fantoorer
Cora
Conu 1D
Banch Norm
Relu
Batch BatchWorm Norm
Conu 1D
Batch Norm
Combination
Module
Self Q,KV
Combination
Module
Conv 1D
Batch Norm
RelU
1D
Batch NothhNorm Norm
RelU
10
Batch Norm
Conv 1D
St-1
Latent Space Navigation
dig
ReiU
Cony 1D
ReLU
Conv 10
Conv
Mg seem dig)
DCnn 1D
RetU
Conv 10
ReULU
Conv 10
Conv 1D
out
My is
1500-9250 R:1
125-11 K3
supervise My
out Mg (S₁ = 0, Zi)
402 110-1139000
150-3125-
1D
FasiFormer
Factformer mer
RetU
ReLU
Zyem
4 1.2256 K-2
1000 100
200-9100
100-9200
256-9256
Combination Module
Positional encoding
Fully connected
Group norm
Conv 1D
SILU
SIUU
P
Group
Figure Implementation dotails of our method when cycle conditioning replaced with the disentangled guidance strategy proposed by (Bhunia al. 2023).
applicd to the embedding layers. A detailed description of the configurations and associated loss functions can be found Figures 3-11. Here, we denote the conditioning sample as C' and the input sample at time step S, distinguish between them.
Results: As shown in Table 6, our method outperforms all competing strategies. It is important to note that none of these methods were originally designed for learning representations from categorical samples, necessitating the inclusion of additional strategies to adapt them for categorical data. We employed the approach outlined in (Lin et al. 2023 Dieleman et al.] 2022). which involves making an extra prediction on the embedding layer before it is fed into the diffusion model.
Additionally, we conducted an experiment where the same embedding input is used at the end of the network to maintain consistency in the intermediale representations, illustrated in Figure 9. However, experiments revealed that this approach encounters significant convergence issues, primarily due to the mutual dependence between the embedding layer and the diffusion model.
For each method, we examine two different configurations, denoted by the sulfixes IE and SE in Table 6. The suffix SE refers to a shared embedding layer that learns to encode both the condition and the imput sample, while IE indicates the use of separate embodding layers for each sample. One the main advantages of an end-to-end embedding layer the ability to learn distinct embeddings for different input samples, which can be achieved with two individual embedding layers (Devlin et al. 2018). Tn this configuration, the network
19

---

Published in Transactions on Machine Learning Research (03/2025)
Regularization term
E|]l E[llzsem 11a]
2 zem
6-2256 8:
Con 230-9236,
236-92566
S
Conv 10
Batch Norm
ReUU
Conv 1D
netU
Batch Norm
Combiretion Module
Self att: Q.K.V
Comination
Modale
Batch Norm
ReiU
Conv 10
Balch Norm
ReiLU
Conv 1D
ReLU
Batch
shared branch
Zsem
shared branch
C
13v27--1226500
Embedding
Fantaomer
Embedding
layer
Σ
wi log 1/
layer
Shared branch
Conv 10
Rell
Conv
Fasti PastFomer orner
ReLUU
FastForme
RelU
FastFormer
ReilU
596. K
24,43
10-2-00
X2-1000
XP-1000
Figure 5: Implementation details for extending the strategy proposed in ReCondSE (Mittal al.
2023 for
categorical data. The embedding layer is supervised with the strategy proposed in (Lin al. 2023 we do not consider the cycle embedding for the entail of this model.
Heart failure prediction
HOur metiod
RCondSE
RCondIE
DiLAESS
DILABIE
InfoDilleson
Mittal al.
2023
Mitsal
2023
Badui al. 2021
DDvln
2018
Devlin al. 2018
heeu
2021b
AUROC AUCPR AUROC AUCPR
AUROC AUCPR
AUROC AUCPR
AUROC
AUCPR
AUROC AUCPR
AUROC
A AUCP
Support Vector 93 80-a.s 19. 14.14.8 91 56.amm 14.78-6%
89.75-0.9
18.34.5%
42.0... 14.41495
90.62-0.44
1493-18
90 10.0.s
14 14.51-4%
91. 91.73405s
15.82495
Gaussian Process
17 91.16.9.4
16.33aax
91.97aas 17.5744%
89. 18.37
13.96.4%
97.17.
15.07.45
80.04cm
13.18 91.07as.8
15.81a6%
Random Forest 00 26-044
13.21.44 88 98.31
11.71aa%
80.41ans
13.06.tth
80.85
14.98.0%
91.47.07
16.84.7%
70.0.s
16.41a7%
90.64.0772
14.44
Ada Boret. 88 76-om 11.96-69
87. 11-0.s 10.59-6%
86.21-051
46-0.57
09 45
08.96-88
87 41-5
17.72-19%
85. 96.104
09 099.99485
GaussionB
3. 14.17.0g 88
11:26a05
87.93.na:
11.34
87.45.0.44
87.37.054
10.50.st
85 G8
09
62.7%
49.4990.2
11.00
Av Kaypa
Ar. AFT
ArBFT
Av. Kapps
Av.1FI
Av. AapKappa AWIEFT Av Aappapp
Av DFI
Av.Kapa
Av.1FI
Ka.Kapss
Ar.
0.882494
0.897
D 86262-6
0.873--9
0.831-995
0.849-45
0.756.m6 0. 08784406 0.857-16
0.741
0.793
0.802-99
D 0.806
0.773
Lung cancer prediction
Support Voctor 00 73....0
U.S1a0%
90.
0.53xck
80.70am2
0.30..6%
00.07.0.17
0.31a0%
80.32.0.
0.48.46
89 05.05a
0.4064%
85 8260.4
0.46.675
Gonseian Procrss:90
94-0-9
0.56-1%
87.57-0:1
0.41-8x
87.81-047
0.42-11%
84-0.31
0.42445
88.27-037
0.43-6
-1-31
0.44-6%
87.6746-0
0.41-58
Random Forest. 92.85.0034
0.71
88.35u0.6
0.44.e8
88.78an.0
0.45..%
89 02.0.99
0.49usm:
90:20.0.3
0.52uet
89.04.0.5
0.47ua
89.38a0.22
0.48.48
Ada Bocet 90. 30.35
0.53,0% 85 65.665
0.36611%
85,91am22
0.36.s%
87 87.05.0.33
0.10a4%
87.43.0.
0.1la78
88 1sss4
0.4324%
85. 162062
0.41.4%
Genesian NB
10-0.50 0.51 88.15-
0.41-x%
88.4
48.43-0a
0.44-1%
74-0.7
0.36.8%
86.52-0m
0.3845%
52-021
0.41-6%
254654
0.43-58
Av. Kappa Av. BFI
API
Av.BFI
Av. Kappa
Av.EF]
Av Kagaa
AF.
Av. Kappa
BFI
Av. Kvggapa
Av.BFI
Av. Ko.Kappu
Av.BFI
1. 1.407% 47 47.44. 1.01as% 16.032asz
0.74xn 41.120% 0.79.e% 41.82.074
0.72.7%
41.25xn.2
0.51a6% 0.802.0%
0.08a1195
43.92un7
Table 6: Performance comparison of our method with liffusion-based representation learning methods
can automatically learn the same embodding for samples that exhibit a high correlation. Our goal is assess whether this type of embedding can improve the performance of our competing strategies, which already benefit from end-to-end learning of the embedding layers.
20

---

Published in Transactions on Machine Learning Research (03/2025)
Regularization term
E E[ l11]
sem
230-9236, I0
6-9256 8:
236-9 211-222666
Wi
Conv 10
Batch Norm
RelU
Conv 10
Norm
neeU
Batch Norm
Combiretion Module
Self att: Q.K.V
Comination
Module
Batch Norm
ReiU
Conv 10
ReiLU
Conv 1D
ReLU
Batch
Batthh
Embedding layer
FassFormer
log P4
K:1
1001->200
1009200
100
C
Conv 10
RetU
FastFormer
ReiU
FaceFormc
ReUU
FastFormer
RelU
yem
Embedding layer
FastFormer
Σ
log
J3N22500
Figure 6: Implementation details for extending the strategy proposed in ReCondIE (Mittal al, 2023 for categorical data
As shown in Table 6, our method achieves improvements of 2.24%, 1.83%, 2.95%, 2.33%, 2.10%, and 2.07% in AUROC compared to the best results of RCondSE, RCondIF, DiS, DifTAFIE, InfoDiffusionSF, and InfoDiffusiinlE for the heart failure prediction task. For lung cancer prediction, our method outperforms the best results of the competing algorithms by 3.24%, 4.03%, 3.66%, 3.53%, 4.05%, and 4.35%. The improvement in lung cancer prediction is notably greater than that in the heart failure prediction, likely because the number of patients with lung cancer in database is significantly smaller than that for heart failure. Consequently, the benefits of the probabilistic representation enabled by cycle conditioning are more pronounced in lung cancer prediction.
In contrast, the competing strategies exhibited similar performance, which may be attributed to our use of the same configuration introduced in (Mittal et al., 2023: Preechakul et al.) 2022 Wang al., (023b) have limited impact on the representations learned by the conditioning encoder. It is possible that different architectures would enhance their effectiveness when applied in their original form, which is specifically designed for inage-related tasks. Additionally, their lower performance on categorical data may stem from the lack of specialized design to extract versatile information from the conditioning encoder.
most of their components. This suggests that the regularization techniques
C Cycle consistency disentangled guidance
We propose our cycle conditioning strategy as a method to enhance the representations learned by diffusion models. Additionally, suggested that this cycle conditioning prevent the diffusion process from
21

---

Published in Transactions on Machine Learning Research (03/2025)
Conv 206-9226 8:
259-92556 8:3
256 256-34, 8:3
1001-223
Fully connected
Batch Norm
Batch Norm
Batch Norm
(x+)
Combination
Self D.K.V
Cominatiio
FastFormer
Re
ReiU
Retu
Conv 10
ReUU
Module
Module
Cone 10
St-1
Embedding layer
Fully cornected
ssem
2xem
y:
w; logy;
4-9256, K.3
400-550 Fully connected
FastFormer
Farthormer
Conv
RelU
ReiU
sem
Embedding layer
2001 1001-93000 9000
100-3300
200-9-200
Fully connected
Wi logy
Grouped 77 chunnes
10x7-42000
4000- 1001
40203 2001
1001
80-9 1000
80-91000 1001
Positional encoding
connected
connected
Fully connected
Fully connected
Fully connected
(t)
7 sem
>
Group
((5em)
Paffine
Zaffine (1 + (Zsem) +
Channel norm
Group norm
Conv 10
(x,))
SIU
SILU
Conv
56 89
35, K33
fifeine fine
Figure 7: Implementation details for extending the strategy proposed in DifAAEIE Prrechakul et al. 2022
for categorical data. The embedding layer is supervisod with the strategy proposed (Lin et al. 2023
overlooking the condition when the hidden condition supervised by extra spelling information of the original lime series. Another strategy to address this issue is disentangled guidance, adapted from the work of (Bhunia et al. 2023). In this approach, the condition is preserved by replacing S, with all-zero tensor and supervising the network to reconstruct the condition from this all-zero tensor combined with the latent condition. This reconstruction is facilitated by the combination modules in the denoising network and is used in conjunction with the conventional loss functions of diffuusin models. The strategy and the definition of the combination module are illustrated Figure
Tables and 8 compare the performance of our method with the alternative strategy of disentangled guidance. As shown, using supervised cycle conditioned network, we achieved 3.53% and 3.00% higher AUROC
22

---

Published in Transactions on Machine Learning Research (03/2025)
206--22 8:
259-92556 8:3
125256 85
256 256-34, 8:3
1001-923
Fully connected
Re
Conv
Batch Norm
ReiUU
Conv
Batch Norm
RetU
Conv 10
Batch Norm
RelU
(x,)
Combination
Module
Self D.K.V
Cominatioo
Module
Cone 10
FastFormer
St-1
Zsem
2xem
shared branch
Zyem
shared branch
S5
C
...
1ax2650
Embedding
Embedding layer
Σ
V:
wi log v;
Conv 10
Cont 1.0
FantFormer
V FastF ormer
ReLU
ReiU
ReiU
ReiUU
4-921, K-1
256-34 8:
1003-92100
100-9200
100-5 200
Grouped by 27 chamnass
4009 1001
00--7000 1003
-000
80-9-0001
1020-91000
60
Positional encoding
connected
Fully connected
Fully connected
Fully Fully connected
(t)
Zgem
Group
f2
Paffine
Laf
(1 + x)f + B
Channel porm
Group norm
((x))
Conv
Conv
4-9256, K3
256-9256,83 K33
Ep
Zaf fine
Combination Module
Figure 8: Implementation details for extending the strategy proposed in DiffAESE PPrecchakul et al. 2022 for categorical data. The embedding layer is supervised with the strategy proposed in (Lin et al. 2023 we do not consider the cycle embedding for the entail of this model.
compared to the best results obtained with disentangled guidance on the two downstream tasks, respectively. Furthermore, the mean AUROC of our method across different downstream classifiers is 2.96% higher than that of the disentangled conditioning.
23

---

Published in Transactions on Machine Learning Research (03/2025)
61339620001
Cony 255-9286, K..3
Conv 43256 K:
256- 9256
256-34
1001-913
Straght-Through Estimator
Wt x
Fully connected
ReLU
Batch Nomm
RetU
Batch Norm
RelU
Conv 10
Batch Norm
RelU
(x,,
Combination Module
Self art: Q.K.V
Comination
Module
Xt-1
Conv 10
FaxtForne
p4
see
Zsem
FazsFormer
4-3256, K:3
17:
wi log vi
Fully connectedd
Fastformer
Fast oomeer
Foot Forfor
ReLU
RelU
RelUU
7em
134+2800
Embedding layer
1002->200
100-9108
100-9100
400-3600
Cycle embedding of diffusion model, according to our experiments this structure has significant. difficulty with convergence.
Figure 9:
This improvement may be attributed to the fact that, in disentangled guidance, the condition is processed differently when provided without the noisy sample compared to when it is combined with the noisy sample. In contrast, our method employs supervision approach where both the input and the denoised sample follow the same pathway within the network. This design helps avoid conflicts between the learned representations in the input of the denoising model and those maintained in its output.
Additionally, disentangled conditioning does not allow for generating multiple representations from a single conditioning sample. Our method addresses this limitation by producing diverse reprosentations from single sample. This capability is especially advantageous when there not enough labeled samples to adequately train downstream classifier.
Moreover, disentangled conditioning introduces extra supervision on an intermediate layer. This supervision involves estimating the elements the spelling vector from an intermediate layer following the combination modules. Such supervision can lead to sparsity that intermediate layer. Since this layer located close the model's output,, it may negatively impact the performance of the denoising model in estimating the noise in its input sample.
Cycle conditioning Disentangied guidance
Cycle conditioning Disentangled guidance
AUROC AUCPR AUROC
AUCPR
AUROC AUCPR AUROC
AUCPR
Support Verlor
91.81100.61
12.49
87.57
19.23-16%
Support Vector
90.51 0.54 86. 0955
0.41
Gaussian Process 90.96-t0.62
13.70+t%%
86. 16.12
08.16-8%
Gaussian Process 88.79-t0.8
0.56+8%
$7.86-a365
0.52-44%
Random Forest
89.48.10.37
14.21
88
0.69
13.61
Random Forest 91.35:t0.48
0.59
88. 0..7
0.48
Adaa Boost 85.12-t0.75
68.62+66%
84.01- 0.43
08. 16-
Ada Boost 86.60-+0.577
0.49+66%
86.92-
0.51+00%
GaussianB
89 80.20 18.4
85.73
L082
19.1%
Ga
$4.39-0.77 0.43-c% 65.62-t0.11
0.46-2%
Table 7: Comparing cycle conditioning and disentangled guidance for heart failure prediction
Table 8: Comparing cycle conditioning and disentangled guidance for lung cancer prediction
D Class-inspired representation learning meets the cycle conditioning
In this section, we propose an end-to-end class-inspired reprosentation learning method that utilizes class labels from a subset of samples develop robust representations for all samples in the database. This approach is particularly valuable for applications where the number of training samples with specific labels is
24

---

Published in Transactions on Machine Learning Research (03/2025)
very limited. Unlike multi-task representation learning, the representations learned by this model focus the characteristics most relevant to predicting particular label.
While this method can eliminate the need for a downstream classifier by integrating the classifier into the representation learning framework, we found that training an additional classifier on the class-specific representations generated by this model is significantly more effective than relying on the model's internal classifier.
For this modeling, we first consider lower bound estimation of our previously proposed method for batch of data that includes at least one sample diagnosed with a specific disease of interest. This can be any kind of disease including heart failure or lung cancer that we already considered as downstream tasks of our model. In this case, the variational bound of our model can be written as follows:
log Po( y) (S1rr Su. y)
log po (S. v) = Eg
log (S ST u)
E,
log po Pp 9)
((S)
pe (1 ST. ST. 4) log (( y)
Eg
log (S +
Essis 129))
(16)
Rg
1ogp(Sr) - log po(p))
+
Σ 111
Rg S.-1.)))
- log Pog +logg(S
(S. 1.))
log ((T)
log (og)o(v) - log po Po: log ((55((1-111)))
T
11
Py()) as ST and y are independent and the distribution of Sy is already known and therefore is not dependent on 0. This equation measures accumulated reconstruction loss over time steps of the denoising process and therefore can be approximated follows:
log 2
121
(17)
v.S,v)
121
For data batches where none of the training samples are diagnosed with the intended disease that aim for the represcutation to Icarn, the information from the labels is treated as latent variable. This approach allows us to estimate the lower bound for these samples using the following formula:
log (S) = Σ Egtv.S.Is..) [ - + (ogg(Sr)- - - log po()) - log pa (S + log Se, S,
1) Bo(s.bs..- [- 1 logp(ST)= T
1 log (o
12
log Peppo (S 3) + log glv. S.
(18)
25

---

Published in Transactions on Machine Learning Research (03/2025)
Using the chain rule, we have S((Se-) (((18(-1))((((y)55-))) Therefore, this equation be reformulated as follows:
Pogpo(S) =
Σ 1)) Ee(s..y.s..1)
1
log logp(Sr) - log pw()
llgg pa (S.-1|S,,)) log (oly, ((15.)))((S(((s
Σ
R Se
Russiss. 1-
1
log ((T)
T logpo(u)
121
log Po y) + log (())())) + (ogg((1((e_1,)))
(19)
Therefore, the equation can be reformulated as:
log (S) = ΣΣ ((u)Sa-1) (- -..( logr(Sr) logro(y)
-og v) + log ((l(S:-1) + log (ogg($)
ΣΣ ( V) + Rs.... (atliisssnna((SSSs...s ()((v))s-1)))
11
(20)
Where E, 9())S(111) (E (1(()))111))) is the entropy of qlyl Sx-1) and therefore measures how spread out the probabilities are in this specific distribution.
From this equation, the distribution (3)(S can formulated as discriminative classifier, however similar to our cycle conditioned strategy it has to implemented in a cycle conditioned manner as it depends to the output the denoising model. Therefore, it is beneficial to reimplement both the cycle and the label conditioning elements with a simple shared common encoder with different heads. In practice we estimate the label y by applying two fully connected layers after (.),,), with 30 unites per each layer.
Our experiments show that the class- cass-insppied cycle conditioning achieves an AUROC of 94.18% the heart failure prediction lask when using Support Vector Machines as the rownstreem classifier. This represents an improvement of 1.01% over the maximum AUROC value obtained with the vanilla version of our cycle conditioning mothod. Additionally, a similar experiment was conducted to prodict heart failure within a four-month time window, using the same data split as described in the main paper.
E Exploring Data Size Variance
In this section, we evaluate how the size of the training data impacts the expressiveness of the representations learned by our diffusion model. We consider heart failure prediction as the downstream task for this experiment. To this end, we utilize two distinct training sets: one from the MIMIC IV database and another from the MIMIC III database, both of which are used to train our representation learning model. After training, we employ these pre-trained models to project both the training and test samples into an embedding space. The test set is selected from the MMMIC IV database, and the data split follows the same procedure outlined in the main paper.
Our training sets from MIMIC-III exclusively contain ICD-9 codes, while our test samples from MIMIC-IV comprise both ICD-9 and TCD-10 codes. To align with the vocabulary of our MIMIC-TV database, index
26

---

Published in Transactions on Machine Learning Research (03/2025)
the tokens from MIMIC-III based their first three characters, matching them with the corresponding tokens in the vocabulary of the MIMIC-IV database.
We compared the performance of our method with SwitchTab (Wu et al. 2024) using the two databases. The results, presented Table 9, demonstrate that our method outperforms this competing strategy, even when the number of training samples significantly reduced. Notably, when trained on only 5% of the training data, our method achieved an accuracy comparable to many competing strategies that were trained on the entire training set of the MIMIC IV database (see Table 7).
IV
MIMIC
AUROC
AUCPR 12. 55%
AUROC
AUCPR
Our method SwichTabb
91.81 0.
81.7 -..52
05. 05.22---6 02993+111%
00.39
09.57-+6%
65.88-F0.88
Table 9: Performance comparison of our method two databases with varying sizes: We employ MIMIC III and MIMIC IV for training our diffusion- representation learning method.
Mutual Information Maximization
Prior Regularization
q (C,, 2sem)
MI Mlc.zsem = B (C.2se)))
log qg))9g (7sem)
R DKL (9o (Zeem (Zsem)
Zeem
200-9224 K-
Conv 1D
Norm
Batch Norm
Conv 10
Batth Norm
(xx)
Combination
Module
Self Q.K.
Comination
Madule
Conv 1D
Conv 1D
Batch Norm
St Com
ReeU
ReLU
ReLU
RelU
ReeU
S
Batch
Batth
shared branch
Z sem
shared branch
Sr
w
13,77-1332660
Embedding layer
FaziFormer
Embedding
Σ wi log 1)
Cone 1D
Conv ID
FastFormer
Fastf-r Fast-ormme armerr
ReiU
ReaU
1001->100
100-9200
Figure 10: Implementation details for extending the strategy proposed InfoDiffusionSE (Wang al. 2023b for categorical data.
27

---

Published in Transactions on Machine Learning Research (03/2025)
Mutual Information Maximization
Prior Regularization
9s (C, 2yem) log q()) (2
Mle.z
DxI. (90 (Zeem Ilp (Zeem)
gem
8-2256 K:3
Batth 255-)226,
Barch Norm
Batch Norm
(x())
Combination
Comination
Banch Norm
Reil
Conv 10
RaLU
Conv 10
ReLU
Module
arts
Module
Conv 1D
Batch
Relu
Conv 1D
RetU
Embedding layer
D:
W; log 1/
4-9256,
256-44 K:
1001-9100 -100
100-1100
Cony 1D
FastFormer
RetU
RelU
ReLU
ReLU
sem
Embedding layer
FactForer
Weg log V.
Figure 11: Implementation details for extending the strategy proposed in InfoDifinsionEE (Wang et al. 2023b for categorical data.
F Notations
In Table 4, Em-MLM denotes continuous data from the embedding layer trained using mask modeling, OutBERT-MLM represents continuous data from one layer before the output of the BERT model trained using mask modeling, and EM-ConAu signifies continuous data from the embedding layer of the BERT model trained using the method proposed in (Chen et al. 2024b).
G Our data transformation strategy versus BERT embedding layer for downstream tasks
In this section, we compare the performance of our proposed data transformation module with that of the BERT embedding layer, which was trained using masked language modeling, on two downstream tasks: heart failure prediction and lung cancer prediction c As shown Tables 10 and 11, our method consistently outperforms the BERT embedding layer across different classifiers. However, the extent of this improvement varies by task, with the performance gain being smaller for lung cancer prediction compared hear failure
'Dttilled results the BERT embedding layer, along with comparisons other classification models, presented in Table
28

---

Published in Transactions on Machine Learning Research (03/2025)
1) Unique characters in the vocahulary
A, B, .z, 0, 1, ,9
Asssining an integer to each character
4, X, ... . 8, ... 0
35,34, ,10, 9, 8
Representing each token by concatenating its integers
H77 'X45'
'H73'
+173
X45'
+ 3345 :
100''
'Y55'
4 Creating multi-hot Spelling vector
1773-th entry
10
1
Figure 12: The vector used to supervise the latent code of the conditioning sample is derived from the spelling of the tokens in w. This spelling information is represented by the methodology illustrated in this figure. The initial direction lowards each spelling index (e.g., 1773) is learned by an element of the dictionary (da), where x represents the index.
prediction. We altribute this difference to the fact that masked language modeling does not capture time- specific nuances in the embedding layer, which may limit its performance, especially in tasks that rely more on time-specific features. When comparing the performance of our embedding layer with that of our model without spelling supervision, we observe that the performance of the diffusion model with cycle conditioning alone is even worse than that of the input to our diffusion model. This is because cycle conditioning without external supervision causes the diffusion model to forget much of the conditioned information, as many of these details change in each iteration of training.
BERT eubbedding layer
Our DT stratogy
BERT embedding layer
Our DT strategy
AUROC
AUCPR
AUROC AUCPR
AUROC
AUCPR
AUROC AUCPR
Support Vector 66. 66.49:2..4
03. 48
76.2110.36
04 11.29.6%
Support Vector 83.47=0.77
0.311%%
$1.122066
0.3811%
02.861.% Gaussian Process 10.
77.33.10.68
06.37 15%
Gaussiau Process 80.14_0.37
1.31
S5.47
0.27.1r%
Random Forest 67.63:t0.44
03.314:4%
78.04:0.63
05, 28.28
Random Forest 84.19=0.61
0.38
85.7340.60
0.3241%
Ada Boast 64.28+t0.47
02.88+69
74.34+00.51
06.72455%
Ada Boost 81.23_0.4r
0.21144%
$2.28
0.37 1.%
NB 60.58.Lo.68
02.
10.04
14..%
GaussianNB 82.12=0.a8
0.29
83.71 0.30
Table 10: Our data transformation strategy versus BERT Table 11: Our data transformation strategy versus BERT embedding layer for heart failure prediction. DT refers to Data Transformation
embedding layer for lung cancer prediction. DYT refers to Data Transformation.
H Computational Complexity
In this section, we provide analysis of the computational complexity of our method compared to its competing strategies. Our method consists of two primary components. The first part is a data transformation strategy that takes an input sample in the form of a one-hot vector of size 1 x W, where W represents the total number of unique tokens in the vocabulary, and transforms it into a contimuous hidden representation of size x m, where m is the size of the embedding layer. The core of this transformation a matrix of size W m. Therefore, the total number of scalar operations required for multiplying this matrix with the input sample is O(W m). For time series with r time steps, the overall complexity becomes O(r W m). In contrast, the BERT embedding layer is lookup table with learnable parameters, where the lookup operation takes O(1) time to retrieve the Icarned embedding for each token in the input sequence. For a time scries of r individual time steps, the computational complexity of BERT's embedding is O(r x 4), where is the number of tokens in each of the r time steps. Since W > q, the computational cost our method is higher than that of the BERT embodding layer, and this cost scales with the size of the vocabulary. However, this comparison assumes that the concatenation of feature embeddings for the entire sequence of tokens in BERT is not prohibitively large and can be directly used as the embedding representation. In cases where the size of
29

---

Published in Transactions on Machine Learning Research (03/2025)
the embedding layer for each token needs to be large enough to capture the necessary contextual information, must introduce an additional linear layer to project the embeddings into lover-dimensional space. This transformation involves multiplying the embeddings with a weight matrix, which results in an additional computational cost of Or x x my x m), where my is the dimension of cach embedding before the lincar projection, and m is the size of the output embedding space afler projection.
The second part of our method is a conditional diffusion model, where the representation learned from the conditioning sample is utilized the final representation of the input sample. Therefore, the computational cost of this part depends heavily the computational cost of the conditioning encoder. Central to conditioning encoder is FastFormer module (Wu et al. 2021). which leverages the additive attention mechanism to reduce the quadratic complexity of Transformers to linear complexity. In Transformers, a sequence of length m requires m x m pairwise computations to account for all combinations of queries, keys, and values. Hence, has O(m²) computational complexity. In contrast, Former avoids the pairwise computation by utilizing additive operations aggregate attention. As a result, the computational complexity required for computing allention reduced to ((n). This represents a significant. advantage in the computational efficiency of our method compared to standard Transformer- based models like SwitchTabb 2024 especially when the size of the embedding layer scales up.
30


Reviews:
[{"review_id": "cIEhy12Q9x", "reviewer": "Reviewer_pmoU", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The method is novel for learning better representations for categorical data using the new powerful generation models - diffusion models.\n2. The paper proposes careful and insightful design components to make sure meaningful representations are learned. \n3. The proposed method shows superior performance.\n4. The exploration of probabilistic representation learning is interesting, further highlighting the strengths of the learning method with diffusion models.\n\nWeaknesses:\n1. The paper has some typos, such as citation format and equations. See the next section.\n2. The paper lacks some reasonable justification for certain design choices. See the next section.\n3. The proposed method is effective on medical data and disease prediction, yet the effectiveness is unknown for other categorical data or downstream tasks.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a new method to learn categorical data representations using diffusion models, aiming to learn more versatile and general-purpose information from categorical data and thus enhancing adaptability to unforeseen tasks. First, the paper proposes a simple data transformation strategy that maps categorical data at each time step to continuous data: a predefined vocabulary is used to construct a multi-hot representation for categorical data at each time step, and then the multi-hot representation is mapped to a fixed size representation through a learnable transformation matrix; the transformation matrix is optimized iteratively. Next, the paper proposes to further learn a compact representation of the transformed continuous data using diffusion models. Specifically, the paper learns a diffusion model conditioned on the compact representation, together with the encoder producing the compact representation. To encourage the extraction of general-purpose representations, the paper further proposes cycle conditioning which intuitively forces the encoder to capture the same content from both the clean and the noisy samples. To enhance the conditioning information, the paper proposes spelling supervision which ensures the latent condition can recover the \"spelling vector\" for each time series. Further, to encourage a denser latent code, latent space navigation is proposed to improve the supervision with spelling vectors. The paper then conducts representation learning on MIMIC III and MIMIC IV datasets and then uses the learned representations (mainly the latent code) to do disease prediction. The learned representations show SOTA classification performance. The paper also demonstrates probabilistic representation learning further enhances the representation, with the help of the stochastic nature of diffusion models. Lastly, ablation studies are conducted to show the effectiveness of cycle conditioning and spelling supervision.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Minor ones:\n1. The authors may consider correcting the citation formats, many citations should be in parentheses.\n2. The authors may consider adding a definition of categorical data.\n3. Section 4.1: in languages, typically sentences are also not of the same length and padding can be used to produce consistent fixed-size vectors for each sentence, so the argument that the embedding layer has \"limitations when dealing with categorical time series\" is not sound. Also, the argument that, training both the embedding layer and a continuous diffusion mode may not be stable, is not convincing - is the argument a conclusion from experiments?\n4. In the paragraph above equation 10, it should be $M_\\theta\\left(g_\\theta\\left(S_t, t, E_\\theta\\left(\\frac{1}{\\sqrt{\\beta_t}} S_t-\\frac{\\sqrt{1-\\beta_t}}{\\sqrt{\\beta_t}} \\epsilon_{t-1}\\right)\\right)\\right)$ instead of $M_\\theta\\left(E_\\theta\\left(S_t, t, E_\\theta\\left(\\frac{1}{\\sqrt{\\beta_t}} S_t-\\frac{\\sqrt{1-\\beta_t}}{\\sqrt{\\beta_t}} \\epsilon_{t-1}\\right)\\right)\\right)$.\n5. Figure 2 caption: it should be \"left panel\" in the last sentence.\n6. Same concern in equation 12 as point 4 above.\n7. Same concern in section 4.2.5 as point 4 above. Moreover, $g_\\theta\\left(S_t, t, E_\\theta\\left(\\frac{1}{\\sqrt{\\beta_t}} S_t-\\frac{\\sqrt{1-\\beta_t}}{\\sqrt{\\beta_t}} \\epsilon_{t-1}\\right)\\right)$ is predicted noise, maybe a justification on why the predicted noise is useful to derive representations.\n\nCritical ones:\n1. Cycle conditioning: whether it is incorporated at all denoising time steps during training? Timesteps near pure noise may not be suitable since the added noise is too much and makes a big difference between clean and noisy data. In this case, incorporating cycle conditioning instead harm the learning of $E_{\\theta}$. Hope the author could provide explanations on choices of time steps and ablation studies.\n2. About the loss function in equation 14: each loss term should have a different magnitude and of a different importance level, I wonder why the authors decided to directly sum all of the loss terms without weighting, and it may be better to have a justification for this.\n3. About the claim of learning \"general-purpose information\" using the proposed method: The downstream task is disease prediction (binary classification), how can this experiment results show learning of \"general-purpose information\"? It would be more convincing if the learned representations could support different types of downstream tasks.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "fnMKHdEd40", "reviewer": "Reviewer_KxR8", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Summary of strengths\n\n- The paper is well-structured, providing a detailed background on related work.\n- The integration of cycle conditioning and spelling supervision into a diffusion-based framework is novel and addresses challenges in representation learning for categorical data. \n- The proposed model outperforms existing state-of-the-art methods and other diffusion-based representation learning techniques.\n- The paper introduces a novel data transformation strategy to capture meaningful information from sequences of categorical tokens, demonstrating superior performance compared to the BERT embedding layer.\n\nSummary of weaknesses\n\n- The evaluation focuses exclusively on the MIMIC datasets, which limits the applicability of the method to other domains like finance or recommendation systems.\n\n- While the paper critiques the high computational demands of graph neural networks for categorical data, it does not provide a computational cost analysis of the proposed diffusion-based method. Given that diffusion models are computationally intensive, a detailed comparison of efficiency and scalability would be valuable.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents an innovative approach to categorical data representation using a diffusion-based method. The proposed solution integrates three core components: diffusion-based representation learning, cycle conditioning, and spelling supervision. These components work systematically to capture meaningful and general-purpose representations from categorical data, addressing key limitations of existing methods. The authors focus on healthcare datasets (heart failure and lung cancer prediction tasks), and demonstrate that their method outperforms state-of-the-art techniques, both in terms of AUROC and AUCPR.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Explicitly present and discuss ablation studies for cycle conditioning and spelling supervision in the main text to underscore their individual contributions to the overall performance, as indicated by the experimental results in Tables 6-7.\n- Provide a detailed computational cost analysis, comparing the proposed method with other state-of-the-art techniques.\n- Consider adding a table to summarize and compare the performance of different diffusion-based strategies from Figures 3–11.\n- Address minor presentation issues:\n  - The term \"dense\" in Figure 1 appears incomplete.\n  - The formatting of tables is inconsistent; some lack top rules, while others lack bottom rules.\n  - Figure 2 is unclear and seems to lack proper reference in the text.", "broader_impact_concerns": "No concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "0c50q23Azi", "reviewer": "Reviewer_FezA", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths:\n\n1) The first component of the proposed framework, which transforms multi-hot representations of tokens into continuous embeddings, is well-motivated and clearly explained. Such design takes the gap between traditional text embeddings and categorical time series data into consideration, while also maintaining robustness against small token modifications. The effectiveness is validated in Section 6, where the authors' method shows superior Spearman correlation compared to embeddings from Masked Language Modeling (MLM), highlighting its ability to capture contextual relationships.\n\n2) The authors incorporate practical strategies to optimize representation learning. For example, cycle conditioning ensures consistency in extracted representations between clean and noisy samples, while spelling supervision guarantees that the learned embeddings retain critical spelling information, thereby preventing the diffusion model(DM) from neglecting the conditioning representation. Furthermore, the authors conduct comprehensive comparative studies against SoTA methods, consistently demonstrating significant improvements in predictive tasks for healthcare applications​.\n\n### Weaknesses:\n\n1) My primary concern lies in the design choice of learning the representations $z_\\text{sem}$ solely as optimal conditioning for the DM, rather than aligning them with its intermediate features. This is because conditioning accounts for a relatively small aspect of the DM, while intermediate features appear to be more critical for reconstruction. While the authors argue that *\"diffusion models do not include a bottleneck layer, so attention has shifted to conditional diffusion models\"*, this is inconsistent with existing UNet-based DMs, which possess a bottleneck layer that could serve as a robust representation space (e.g. https://arxiv.org/abs/2112.03126, https://arxiv.org/abs/2303.09769). Similarly, features from the middle blocks of transformer-based DMs are highly informative, even without an explicit bottleneck (e.g. https://arxiv.org/abs/2401.14404, https://arxiv.org/abs/2410.06940). Though Table 5 shows the superiority of the condition-based method towards those that use unconditional DMs, its effectiveness in more general domains, such as vision or language tasks, remains unverified.\n\n2) Due to the training objective discussed above, the authors need to introduce additional ad-hoc mechanisms, such as cycle conditioning to ensure robustness to noise (a property naturally achieved by intermediate features) and spelling supervision with latent navigation to sufficiently train the conditioning representations. Other approaches may not require such modifications. Also, the ablation study in Table 9 indicates that cycle conditioning contributes much less than spelling supervision, raising questions about the necessity and effectiveness of such a design.\n\n3) The iterative construction of the mapping in the data transformation step, while intuitive, introduces potential computational complexity, as it involves inverting a $d\\times d$ matrix at each epoch, where $d$ represents the total number of tokens.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In this work, the authors present a diffusion-based representation learning framework for categorical (discrete) time series data. Their approach incorporates two key components: (1) a data transformation module that maps the multi-hot representations of tokens at each time step into continuous vectors, while preserving token set similarities, and (2) a diffusion-based representation learning component designed to reconstruct the original sequence and extract robust features as conditioning embeddings. To enhance representation learning, the framework introduces three regularization techniques—cycle conditioning to ensure consistency in extracted features, spelling supervision to maintain conditioning influence during denoising, and latent space navigation to address sparsity in spelling information by introducing learnable directional codes.\n\nSuch a framework significantly improves over SoTA methods, particularly in predictive healthcare tasks on MIMIC-III and MIMIC-IV datasets.​", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1) More ablation studies: \n(1) An ablation study on the data transformation module to compare the performance of the proposed embeddings against MLM-based embeddings for downstream tasks. (2) A comparison between cycle conditioning and explicit alignment of $E_\\theta(S_0)$ and $E_\\theta(\\hat{S}_0(S_t))$ (e.g. MSE loss between clean & noisy representations) to assess the necessity of cycle conditioning.\n\n2) Writing improvements: \n(1) Include a brief preliminary explanation of time series with categorical data to help readers understand the context better. Clarify the distinction between \"time step\" as used in time series analysis and DMs. (2) Provide relevant citations for the second part of the method to better contextualize its novelty and highlight connections to existing work. Without this, the method may appear overly ad-hoc and harder to follow.", "broader_impact_concerns": "No ethical concerns", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper introduces a novel diffusion-based approach for categorical data representation. However, its evaluation is limited to the MIMIC datasets, raising concerns about the method’s generalizability to other domains.
Decision: Accept (poster)