Reproducing codebook training: complete codebook collapse and missing preprocessing details

[jarch100]

Reproduction question

Thank you for releasing GS-Quant. I am trying to reproduce the end-to-end pipeline for FB15K-237 and WN18RR, but the RQ-VAE/codebook stage collapses to one (or two) active codes per layer.

I would appreciate guidance on the exact preprocessing/checkpoint setup or release artifacts needed to reproduce the reported entropy curves.

What I used

• Raw datasets and the required DIFT / KG-Fit-related resources prepared locally.
• Complete entity_info_llm_hier.json for both datasets.
• Local text embeddings generated from mixedbread-ai/mxbai-embed-large-v1 (cached revision b33106f585b9ce46904ad7443a3b52b7a63e231c), with entity labels and LLM descriptions encoded separately and concatenated.
• Graph embedding stage: RotatE, hierarchy type llm, hidden dim 1024, batch size 512, learning rate 1e-4, rho=0.6.
• Codebook architecture/settings from the paper: embedding dim 512, 4 codebook layers, codebook size 1024, encoder [512, 512, 512], reconstruction layers/heads 2/4, parent recon count 5, dropout 0, LR 1e-4, commitment weight 0.25.
• Dataset-specific settings: FB15K-237 batch size 14541 / 500 updates; WN18RR batch size 16348 / 1000 updates.

The data checks passed: entity counts, ID ordering, dimensions, and finite values all match the expected shapes.

Observed behavior

Using the released codebook implementation:

• FB15K-237: full-dataset entropy becomes 0.0000 at the first evaluation, with 1 active code out of 1024 in each layer.
• WN18RR: same result (0.0000, 1 active code/layer).

I also tested a controlled local variant that:

1. uses the full entity set for entropy measurement;
2. keeps LR constant at 1e-4;
3. makes max_steps mean actual optimizer updates; and
4. changes the neighbor-loss average from sum(layer_losses) / (len(layer_losses) + 1e6) to an ordinary average, while safely handling entities with no neighbors.

The neighbor term then became non-negligible (~0.72 instead of ~1e-6), but WN18RR still collapsed completely. FB15K-237 reached only 2 active codes/layer (mean entropy ~0.67) and became unstable.

Clarifications / artifacts requested

Could you please share or confirm:

1. The exact Git commit used for the paper experiments.
2. The exact PLM/model revision used for textual entity embeddings, including text templates, pooling, normalization, and the fusion coefficient rho.
3. Whether pretrained entity_init_embeddings.npy, graph-embedding checkpoints, generated descriptions, or codebook checkpoints are available.
4. The intended codebook initialization, optimizer/scheduler, validation/entropy protocol, and handling of empty neighbor sets/dead codes.
5. Whether the neighbor-loss denominator in the released code is intentional.
6. A complete command/config for reproducing Table 3 / Figure 6 entropy trajectories.

I can provide logs and the minimal patch used for the controlled experiment if helpful. Thank you!

[mikumifa]

We have provided the codebook checkpoints from our previous training runs, and we hope they will be helpful.

config.json

entity_quantized.json

tokens.json

train.log

[jarch100]

Thank you very much for your helpful reply and for sharing the files. Sorry to bother you again.

We inspected the shared config.json, train.log, entity_quantized.json, and tokens.json. They were very useful: they show that the successful run used a different RotatE checkpoint path and codebook behavior than the current public source. With our locally generated embeddings, the codebook still collapses to one code per layer.

Would it be possible to share the following, for both FB15K-237 and WN18RR, if you are able to?

1. The complete processed_data/<dataset>/checkpoints/RotatE/ directory used for your experiments—especially every .npy file (for example, entity_embedding.npy and any associated intermediate/input embeddings).
2. The complete codebook/ folder from the exact version used to produce the shared train.log and entity_quantized.json, including any scripts, configuration files, or preprocessing utilities not present in the public repository.

These would let us reproduce the pipeline faithfully and isolate whether the remaining difference comes from the embeddings, the codebook implementation, or preprocessing. Thank you again for your time and help.

---

非常感谢您的回复，也非常感谢您分享这些文件。再次打扰您，十分抱歉。

我们已经检查了您分享的 config.json、train.log、entity_quantized.json 和 tokens.json。这些文件非常有帮助：它们表明成功运行使用了与当前公开代码不同的 RotatE checkpoint 路径和 codebook 行为。使用我们本地生成的 embeddings 时，codebook 仍然会坍塌为每一层只使用一个 code。

如果方便的话，您是否可以为 FB15K-237 和 WN18RR 两个数据集分享以下内容？

1. 您实验中使用的完整 processed_data/<dataset>/checkpoints/RotatE/ 目录，尤其是其中所有的 .npy 文件（例如 entity_embedding.npy，以及任何相关的中间或输入 embeddings）。
2. 用于生成您分享的 train.log 和 entity_quantized.json 的完整 codebook/ 文件夹，即当时使用的准确版本，包括公开仓库中未包含的脚本、配置文件或预处理工具。

这些内容将帮助我们忠实复现整个流程，并判断剩余差异究竟来自 embeddings、codebook 实现，还是预处理步骤。再次感谢您的时间和帮助！