{"id":"5c877d76-6310-4a82-9e04-041be713eec5","arxiv_id":"2412.13503","paper_version":2,"verdict":"CONDITIONAL","confidence":"HIGH","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":4,"one_line_summary":"A VAE with a diffusion prior synthesizes entity-pair embeddings for underrepresented relations, improving document-level relation extraction on long-tail relations.","lead":"This paper proposes a data augmentation method that generates new entity-pair embeddings with a VAE combined with a diffusion model, and uses them to train a document-level relation extraction model on rare relations. It reports small but statistically significant F1 gains on Re-DocRED and DWIE, especially for uncommon relations.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"No evidence that the Stage 2 generative model remains valid for the retrained Stage 3 encoder; gains may stem from generic perturbation rather than faithful relation-wise oversampling.","rationale":"The reader's weakest_assumption identifies exactly the concern I consider most load-bearing: the distribution shift between the Stage 1 encoder used to train VaeDiff and the Stage 3 encoder that receives its pseudo-features. This concern is not merely theoretical. The paper's own ablation shows that unstructured Gaussian noise provides a meaningful fraction of the improvement, so the specific claim that learned relation-wise distributions drive the gains hinges on the pseudo-features being valid in the Stage 3 embedding space. Without evidence of stability, the improvement could be attributed to a generic augmentation/regularization effect rather than to the proposed generative mechanism. The appendix derivation error (Eqs. 31-32) is a real correctness issue, but it is secondary: the empirical objective in Eq. 9 can be viewed as a heuristic regardless of whether the ELBO derivation is clean. The algorithm's label-handling ambiguity is also secondary because the loss computation, while under-specified, is implementable. The distribution-shift issue, by contrast, threatens the interpretation of the main experimental results. I agree with the reader's conditional verdict: the paper should be accepted only if the authors either demonstrate that the embedding shift is negligible (e.g., via reconstruction-error monitoring or an oracle fine-tuning experiment) or explicitly reframe the contribution as a regularization technique rather than faithful oversampling. The concrete test above would settle this with modest additional computation.","tokens_in":16864,"tokens_out":5370,"duration_ms":53924,"concrete_test":"Instrument Stage 3 as follows: at the warmup boundary (epoch n) and again at the end of training (epoch 30), extract positive entity-pair features from the current Stage 3 encoder for a fixed held-out set, run them through the frozen VaeDiff encoder, and compute the mean squared reconstruction error (the first term of Eq. 4). If the reconstruction error at epoch 30 is substantially larger than at epoch n (e.g., more than 20% relative increase), the generative model is operating on a shifted embedding distribution, confirming that pseudo-features become stale. As a complementary check, run an oracle variant in which VaeDiff is fine-tuned for a few epochs on Stage 3 encoder features at the start of augmentation; if this oracle's F1 exceeds the published result by more than the significance margin, the frozen-VaeDiff design is leaving performance on the table and the shift is material.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The central claim is that VaeDiff-DocRE improves long-tail DocRE by generating faithful entity-pair representations for minority relations from learned relation-wise distributions. This requires the VaeDiff module, trained on Stage 1 encoder representations (Section 2.3.2), to produce pseudo-features that lie in the same embedding space as the Stage 3 encoder's representations. However, Stage 3 retrains the baseline DocRE model from its initial weights, and augmentation begins after only n warmup epochs (5 for Re-DocRED, 10 for DWIE, Section 3.1). At that point, the Stage 3 encoder is far from the Stage 1 final encoder, and it continues to drift during training. The paper provides no analysis, adaptation, or feature-alignment mechanism for this distribution shift. The ablation in Table 4 shows that simply adding Gaussian noise to positive-pair representations (Gaus Noise) already improves over KD-DocRE by 0.13 F1, while the full VaeDiff adds 0.45 more. This leaves open the possibility that the incremental gain comes not from faithful oversampling of minority relations but from a structured but stale feature distribution that acts as a regularizer. If the pseudo-features are out-of-distribution with respect to the Stage 3 encoder, the method is not performing data augmentation in the claimed sense, and the interpretation of the SOTA improvements as 'effectively addressing the long-tail distribution problem' is unsupported. This is the most load-bearing gap because it directly connects the proposed mechanism to the reported empirical gains.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes VaeDiff-DocRE, a three-stage data augmentation framework for document-level relation extraction (DocRE). In Stage 1, a baseline DocRE model is trained to produce entity-pair representations that form relation-wise clusters. In Stage 2, a variational autoencoder (EP-VAE) is trained on these representations, with a conditional diffusion model serving as the latent prior (the 'VaeDiff' module). In Stage 3, the baseline DocRE model is retrained from scratch, and after a short warm-up the frozen VaeDiff generates pseudo entity-pair representations for positive pairs, which are mixed with the model's own features during training. Experiments on Re-DocRED and DWIE report consistent improvements over prior state-of-the-art methods, and ablations show that the augmentation contributes beyond simple Gaussian noise. The paper also releases code and reports average results over ten runs with significance tests.","tokens_in":17123,"tokens_out":10013,"duration_ms":88529,"significance":"If the reported improvements are robust and the augmentation mechanism works as claimed, the main contribution is a practical embedding-space augmentation method for long-tail DocRE, and it would be, to the authors' knowledge, the first application of diffusion probabilistic models to DocRE. The strengths are the reproducible experimental setup (two benchmarks, ten runs, significance tests, released code) and the clear isolation of the augmentation effect through the No Aug and Gaus Noise ablations. However, the central 'faithful oversampling' interpretation is currently not supported because of the distribution shift between the Stage 1 encoder used to train VaeDiff and the retrained Stage 3 encoder; the incremental gain over Gaussian noise could be a regularization effect rather than relation-wise oversampling. This distinction is important for the paper's claim to 'effectively address the long-tail distribution problem.'","major_comments":[{"comment":"The VaeDiff module is trained on entity-pair representations extracted from the Stage 1 encoder, but Stage 3 retrains the baseline DocRE model from scratch and starts augmentation after only n warmup epochs (n=5 for Re-DocRED, n=10 for DWIE, Section 3.1). No mechanism aligns the Stage 3 encoder's embedding space with the Stage 1 space, so the pseudo-features can become stale or out-of-distribution as the Stage 3 encoder drifts. The t-SNE visualization in Section 3.5 does not specify which encoder produced the 'encoded data', and it provides no quantitative distribution-distance measure. The Gaus Noise ablation in Table 4 shows that generic perturbation already improves F1 by 0.13 over KD-DocRE, while the full VaeDiff adds 0.45 more; this leaves open the possibility that the additional gain reflects a structured but stale feature distribution acting as a regularizer rather than faithful oversampling of minority relations. Please add a quantitative analysis of the distribution mismatch across Stage 3 training (e.g., MMD or centroid distances between generated and encoded features), or adapt the generator during Stage 3, or initialize Stage 3 from the Stage 1 encoder.","section":"Sections 2.3.2 and 2.3.3, Algorithm 1"},{"comment":"The derivation of Eq. (9) is invalid as written. Eq. (31) is a valid lower bound on the ELBO in Eq. (30) after replacing the intractable log p_psi(z0) with the diffusion ELBO, but the displayed chain in Eq. (32) states that this lower bound is at most the original ELBO and at most E[log f_theta], which drops the KL terms and does not establish the claimed 'valid ELBO formulation'. The text should be rewritten to say that Eq. (31) is a lower bound that is maximized as a surrogate for the original ELBO, and that the regression objective L_Diff is then used as a further surrogate for the diffusion ELBO term. Additionally, the expectation in the last term of Eq. (9) should be over the forward diffusion process q(z_t | z_0), not merely over q_phi(z_p | p).","section":"Appendix A.4, Eq. (31)-(32) and Eq. (9)"},{"comment":"The conditioning in the Diffusion Prior is the sum of the label embeddings of all ground-truth relations of an entity pair, and in Stage 3 the generator produces pseudo-features for whole positive pairs rather than for individual relations. The paper describes the modeling objects as relation-wise distributions (Section 2.2), but the trained distribution is actually over entity-pair representations conditioned on multi-label vectors. For a pair that expresses both head and tail relations (e.g., P580 with P582, as discussed in Section 3.5), the conditional signal is a mixture, so the generated feature does not specifically target the minority relation. To support the long-tail claim, the authors should either generate per-relation pseudo-features (conditioning on a single label and assigning that label to the generated instance) or analyze how often the augmented pairs are ones whose minority relations are not accompanied by majority relations.","section":"Sections 2.2.2, Eq. (7), and Algorithm 1"},{"comment":"The main comparison to KD-DocRE is confounded by the baseline loss change: Stage 1 replaces KD-DocRE's Adaptive Focal Loss with PMTEM loss plus Supervised Contrastive Learning (Section 2.3.1), so the F1 gains over KD-DocRE in Table 1 cannot be attributed solely to the augmentation module. The No Aug row in Table 4 controls for the loss and is weaker than KD-DocRE (75.20 vs. 75.31), which does show that the final gain over KD-DocRE requires augmentation; nevertheless, the paper should either report the same baseline with and without augmentation (i.e., KD-DocRE with PMTEM+SCL but no VaeDiff) or explicitly acknowledge that the reported SOTA comparison includes the loss change as a co-variate.","section":"Sections 2.3.1 and 3.3, Table 1"}],"minor_comments":[{"comment":"The text states that DWIE has 9 out of 66 relation labels in the long tail, but Table 3 lists 65 relations; please clarify the correct count.","section":"Section 1 and Table 3"},{"comment":"The caption of Figure 5 is missing the marker symbol for the encoded data; it says 'represents the encoded actual entity pair representations' but does not show which symbol (e.g., a dot) is used.","section":"Section 3.5 and Figure 5"},{"comment":"The phrase 'Sample sampling' in line 23 of Algorithm 1 should be 'Sample noise', and the comment in lines 11-18 would be more readable if it explicitly stated that self-conditioning uses the previous estimate with stop-gradient.","section":"Algorithm 1"},{"comment":"The acronym 'LT F1' appears in the description of Re-DocRED metrics; it should be 'LTail F1' for consistency with Tables 1 and 4.","section":"Section 3.1, Metrics"},{"comment":"The paper reports that all p-values are less than 0.05 following Dror et al. (2018), but it does not state which test was used (e.g., paired bootstrap or approximate randomization); please specify the test and the number of bootstrap samples.","section":"Section 3.1, Implementation details"},{"comment":"The Limitation section mentions longer GPU training and unverified specialized domains, but it does not mention the distribution-shift issue between the training of VaeDiff and its deployment in Stage 3; a sentence acknowledging this would align the stated limitations with the actual risks of the method.","section":"Limitation section"}],"recommendation":"major_revision","confidential_remarks":"The code release, ten-run averages, and significance tests are commendable and make the empirical claims easy to check. The main gate is the distribution-shift concern in Stage 3 and the invalid derivation of Eq. (9); both are fixable with additional experiments and a careful rewrite. The 'end-to-end' framing in the title and abstract is overstated given that the generator is frozen during Stage 3, and the multi-label conditioning in Eq. (7) weakens the 'relation-wise' augmentation claim. The paper's novelty claim of applying a diffusion prior to DocRE is plausible but should be checked against any concurrent work in the area."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"The paper is a serious attempt at embedding-space augmentation for long-tail DocRE, and the core combination is new: a VAE for relation-wise distributions with a diffusion prior to handle multi-label, wrapped in a three-stage training scheme. That's a clean idea, and the experiments are above average for the field — ten runs, significance tests, ablations, released code. The gains are modest but consistent, especially on LTail F1, and the ablation shows the full VaeDiff beats both Gaussian noise and no augmentation. Credit where it's due: the empirical work is reproducible and the central claim, that this helps long-tail relations without hurting common ones, is supported on two benchmarks.\n\nThe soft spots are real but fixable. First, the derivation of Eq. (9) in Appendix A.4 is wrong. Equation (32) drops the KL terms in an inequality that doesn't go through, so the theoretical justification as written is invalid. This doesn't necessarily kill the method — the loss still works empirically — but the paper needs to correct the math or remove the overclaim. Second, Algorithm 1 is ambiguous about labels: the pseudo-features are generated for positive entity pairs, but the loss is computed against the full label vector y on the concatenated set. The authors need to spell out how generated features are assigned labels. Third, the distribution-shift concern is legitimate but the stress-test overstates it. There is qualitative evidence in Figure 5 that the generated and encoded features overlap in the final model, so it's not true that there is 'no evidence.' What is missing is a quantitative check of whether the pseudo-features stay in-distribution during the early epochs of Stage 3, when the encoder is still drifting. Given the Gaussian-noise ablation already gains 0.13 F1, a skeptic can reasonably ask how much of the full gain is faithful oversampling versus structured regularization. That's a fair question for the authors to address, but it's a request for more analysis, not a demonstrated fatal flaw.\n\nOverall, I'd send this to peer review. The method is novel, the experiments are solid, and the issues are presentation and analysis gaps rather than wrong results. A serious referee would want the derivation corrected, the label handling clarified, and a distribution-shift analysis added. It's a long-tail DocRE paper, so the audience is specialized, but within that audience it's worth the referee time.","headline":"A credible but modest embedding-space augmentation recipe for long-tail DocRE, with a wrong appendix derivation and an unresolved distribution-shift question; deserves review, needs revision.","tokens_in":17682,"tokens_out":2567,"would_cite":false,"duration_ms":25350,"reading_group":"maybe","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"The paper claims that VaeDiff-DocRE, a data-augmentation framework that generates synthetic entity-pair representations in embedding space using a VAE with a diffusion-model prior, outperforms existing DocRE methods on long-tail relations…","keywords":["document-level relation extraction","data augmentation","long-tail distribution","variational autoencoder","diffusion probabilistic model","embedding space augmentation","multi-label classification","class imbalance"],"falsifier":"Re-run stage 3 on Re-DocRED with the stage-2 generative model left frozen but the document encoder re-initialized; if the reported gains in LTail F1 and overall F1 vanish or reverse, the augmentation depends on the original embedding space rather than on faithful oversampling, and measuring the distance between generated pseudo-features and real stage-3 features for the same relations would confirm whether the generator is sampling from a stale distribution.","tokens_in":16640,"feed_emoji":"🧠","tokens_out":9806,"duration_ms":81092,"temperature":0.7,"pith_summary":"Document-level relation extraction (DocRE) identifies relations between entity pairs across a whole document, and real datasets are heavily imbalanced: most pairs have no relation and a few relations cover most positive examples. This paper claims that the imbalance can be treated directly by generating extra training examples for underrepresented relations in the embedding space, rather than by redesigning loss functions that often trade minority-class gains against majority-class losses. It proposes VaeDiff-DocRE, which trains a variational autoencoder on entity-pair representations, uses a diffusion model as the VAE's latent prior so that generation can be conditioned on multiple relation labels at once, and mixes the resulting pseudo-features into the training batches of a retrained DocRE model. On Re-DocRED and DWIE, the framework reports higher F1 than prior state-of-the-art DocRE methods, with the largest gains on the long-tail F1 metric for rare relations. The point of the work, if correct, is that embedding-space oversampling with a learned relation-wise generative model is a practical, backbone-agnostic cure for long-tail DocRE.","feed_headline":"Diffusion-backed VAE lifts rare relation extraction","feed_subtitle":"A three-stage framework synthesizes entity-pair features for rare relations, raising F1 on both benchmark datasets.","key_machinery":"The load-bearing mechanism is the VaeDiff augmentation module, made of two parts. The first is EP-VAE, a variational autoencoder whose encoder projects entity-pair representations into a Gaussian latent space and whose decoder reconstructs them; the second is the Diffusion Prior, a denoising diffusion probabilistic model that replaces the Gaussian prior and is trained to predict the clean latent vector from a noised version. The diffusion model is conditioned through classifier-free guidance on a relation-label embedding computed as the sum of embeddings of all positive labels, which is what allows multi-label generation. The module is inserted into a three-stage hierarchical loop: stage 1 trains a baseline DocRE encoder and classifier, stage 2 trains VaeDiff on that encoder's entity-pair features, and stage 3 retrains the DocRE model while combining its own features with pseudo-features sampled from VaeDiff. The VAE supplies a compact, reconstructable representation space; the diffusion prior supplies a richer, class-conditioned latent distribution than a Gaussian; and the three-stage loop is what lets the augmentation adapt to the classifier being trained.","core_discovery":"The central discovery, stated on the paper's own terms, is that the representations of entity pairs produced by a trained DocRE encoder organize into distinct relation-wise clusters, and that these clusters can be captured and sampled from. VaeDiff-DocRE builds an Entity-Pair VAE (EP-VAE) that reconstructs entity-pair representations from a compact latent space, then replaces the usual Gaussian prior on the latent variable with a denoising diffusion model. The diffusion model is class-conditioned on the sum of the embeddings of all positive relation labels for a pair, which lets one generated pseudo-feature encode several co-occurring relations. A three-stage training loop first trains a baseline DocRE model, then trains VaeDiff on the baseline's entity-pair representations, and finally retrains the DocRE model with the generated pseudo-features mixed into each batch. The authors report that this beats the compared methods on both benchmarks, for example by 1.6 and 2.27 F1 points on the DWIE development and test sets, and improves LTail F1 on Re-DocRED while keeping majority-relation F1 competitive.","pith_inferences":["A direct test of the framework's logic would freeze the stage-2 generative model while re-initializing the stage-3 encoder; if the gains collapse, the pseudo-features are tied to the stage-1 embedding geometry and the improvement is at least partly a regularizing effect rather than faithful oversampling.","A cheaper baseline worth testing is a conditional VAE with a plain Gaussian prior (or tuned Gaussian noise on real features); comparing against it would isolate what the diffusion prior contributes beyond added stochasticity.","The same embedding-space augmentation recipe could transfer to other extreme multi-label document tasks, such as event extraction or fine-grained entity typing, wherever a trained encoder yields label-conditioned clusters."],"forward_implications":["Rare relations receive synthetic training examples in embedding space without new annotation, directly addressing positive-positive imbalance instead of re-weighting loss terms that trade majority-class accuracy for minority-class gains.","Because the stage-1 and stage-3 models can be any DocRE backbone, the augmentation module is portable: upgrading the underlying encoder or classifier should preserve the benefit.","Conditioning generation on the full multi-label vector means a single pseudo-entity-pair can carry several co-occurring relations, matching the multi-label structure of DocRE.","On the reported experiments, the framework raises overall F1 by 0.58 (BERT) and 0.68 (RoBERTa) on the Re-DocRED development set and by 2.27 F1 points on the DWIE test set, while improving LTail F1 on Re-DocRED test by 0.54 and 0.42 points over the state of the art."],"supporting_citations":[{"why":"This work supplies the VAE formulation, the ELBO objective, and the reparameterization trick that EP-VAE relies on.","marker":"(Kingma and Welling, 2013)"},{"why":"This work supplies the denoising diffusion probabilistic model and the training objective that the Diffusion Prior extends and modifies.","marker":"(Ho et al., 2020)"},{"why":"This work supplies classifier-free guidance, used to condition generation on relation-label embeddings.","marker":"(Ho and Salimans, 2021)"},{"why":"This work supplies the equivalence between the VAE cross-entropy term and the diffusion regression objective that the co-training loss in Eq. (9) is built on.","marker":"(Wehenkel and Louppe, 2021)"},{"why":"This work provides the KD-DocRE backbone and axial-attention entity-pair representation structure that the stage-1 baseline adopts.","marker":"(Tan et al., 2022a)"},{"why":"This work provides the Re-DocRED dataset and the Freq F1 and LTail F1 metrics used to measure long-tail performance.","marker":"(Tan et al., 2022b)"},{"why":"This work supplies the PMTEM loss and supervised contrastive loss used in the baseline objective of stage 1.","marker":"(Guo et al., 2023)"},{"why":"This work provides the DWIE dataset used as the second benchmark.","marker":"(Zaporojets et al., 2021)"},{"why":"This work supplies the self-conditioning technique used to improve the Diffusion Prior's denoising predictions.","marker":"(Chen et al., 2023)"}],"fun_headline_variants":["Diffusion-conditioned VAE augments rare DocRE relations","Generative pseudo-features fix long-tail relation extraction","EP-VAE with diffusion prior lifts DocRE performance","Hierarchical VAE augments DocRE entity-pair space","Sampling rare relation clusters with diffusion VAE"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The load-bearing premise is that the stage-2 generative model, trained on representations from the stage-1 encoder, still produces useful pseudo-features after the stage-3 retraining changes the encoder's embedding space; the paper does not demonstrate that this transfer holds.","fun_headline_variants_meta":{"raw":{"variants":["Diffusion-conditioned VAE augments rare DocRE relations","Generative pseudo-features fix long-tail relation extraction","EP-VAE with diffusion prior lifts DocRE performance","Hierarchical VAE augments DocRE entity-pair space","Sampling rare relation clusters with diffusion VAE"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000235,"raw_usage":{"total_tokens":1485,"prompt_tokens":918,"completion_tokens":567,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":534,"completion_tokens_details":{"reasoning_tokens":488}},"tokens_in":534,"tokens_out":567,"duration_ms":5699,"temperature":1.0,"reasoning_tokens":488,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-11T13:04:11.801992+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Re-run stage 3 on Re-DocRED with the stage-2 generative model left frozen but the document encoder re-initialized; if the reported gains in LTail F1 and overall F1 vanish or reverse, the augmentation depends on the original embedding space rather than on faithful oversampling, and measuring the distance between generated pseudo-features and real stage-3 features for the same relations would confirm whether the generator is sampling from a stale distribution.","supporting_citations":[{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"This work supplies classifier-free guidance, used to condition generation on relation-label embeddings."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"This work supplies the equivalence between the VAE cross-entropy term and the diffusion regression objective that the co-training loss in Eq. (9) is built on."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"This work supplies the PMTEM loss and supervised contrastive loss used in the baseline objective of stage 1."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"This work provides the DWIE dataset used as the second benchmark."}],"review_version":1}