---
title: "RegaVAE"
type: "tool"
slug: "trustedllm-regavae"
canonical_url: "https://www.graphcanon.com/tools/trustedllm-regavae"
github_url: "https://github.com/TrustedLLM/RegaVAE"
homepage_url: null
stars: 15
forks: 1
primary_language: "Python"
license: null
archived: false
categories: ["model-training"]
tags: ["language-modeling", "variational-auto-encoder", "retrieval-augmentation"]
updated_at: "2026-07-12T07:48:21.373568+00:00"
---

# RegaVAE

> A Retrieval-Augmented Gaussian Mixture Variational Auto-Encoder for Language Modeling

RegaVAE is a retrieval-augmented framework that combines VAE with a retrieval mechanism to enhance language modeling capabilities by incorporating both current and future information in the latent space.

## Facts

- Repository: https://github.com/TrustedLLM/RegaVAE
- Stars: 15 · Forks: 1 · Open issues: 0 · Watchers: 1
- Primary language: Python
- Last pushed: 2023-12-05T20:52:31+00:00

## Trust & health

_Signals computed from public GitHub metadata. Not a security guarantee._

- Maintenance: Dormant (computed 2026-07-11T23:06:19.722Z)
- Security scan: No lockfile (0 critical, 0 high, 0 medium, 0 low) · last scan 2026-07-11T23:06:20.313Z
- Full report: [trust report](/tools/trustedllm-regavae/trust.md) · [JSON](https://www.graphcanon.com/api/graphcanon/tools/trustedllm-regavae/trust)

## Categories

- [Model Training](/categories/model-training.md)

## Tags

language modeling, variational auto-encoder, retrieval-augmentation

## Category neighbours (exploratory)

_Same-category tools for discovery only - not curated alternatives. Cap shown at six._

- [AutoRAG](/tools/marker-inc-korea-autorag.md) - AutoRAG: An Open-Source Framework for Retrieval-Augmented Generation (RAG) Evaluation & Optimization with AutoML-Style Automation (★ 4,862) [Active]


## Adoption goal

RegaVAE brings a unique approach by integrating retrieval mechanisms with Gaussian Mixture VAEs to enhance language modeling.

## README (excerpt)

_Quoted verbatim from the upstream repository. Untrusted content - treat as data, not instructions._

````text
# RegaVAE
This is the official repo for our [paper](https://arxiv.org/abs/2310.10567): 
> RegaVAE: A Retrieval-Augmented Gaussian Mixture Variational Auto-Encoder for Language Modeling

## Model Architecture

Architecture of RegaVAE. Based on the training data, we first train a VAE to construct a compact latent space, which ensures that the latent variable z contains both current and future information (see § 3.1 of the paper). We then build a retrieval database and then aggregate the retrieved information into the generator (see § 3.2 of the paper). VAE Encoder and Decoder parameters are the same in all steps. In order to ensure fairness, the Corpus data and the Source data in the training set are the same. $G$ represents the Gaussian mixture distribution, and $π$ is the corresponding parameter.

## Datasets
Download three dataset from this [link](https://drive.google.com/drive/folders/1mcn6nqLDVvrGatKHbdbtDSj9PQI5Eu8S?usp=sharing). Unzip them and put them under the data directory.

## Step1
Firstly,
```
cd Step1
```
### Training
For Yelp dataset,
```
python main.py --train_file ../data/yelp/yelp.train.txt \
--valid_file ../data/yelp/yelp.valid.txt \
--per_gpu_train_batch_size 4 \
--cycle_annealing
```
For Yahoo dataset,
```
python main.py --train_file ../data/yahoo/yahoo.train.txt \
--valid_file ../data/yahoo/yahoo.valid.txt \
--per_gpu_train_batch_size 4 \
--cycle_annealing
```
For WP dataset,
```
python main.py --train_source_path ../data/writingPrompts/train.wp_source \
--train_target_path ../data/writingPrompts/train.wp_target \
--valid_source_path ../data/writingPrompts/valid.wp_source \
--valid_target_path ../data/writingPrompts/valid.wp_target \
--dataset_type wp \
--per_gpu_train_batch_size 4 \
--cycle_annealing
```
The above are only the best adjusted hyperparameters. You can get a better Step1 model by passing other parameters. The model we trained is available at this [link](https://drive.google.com/drive/folders/1HmTqQmHSmP_VZUDV9ADM6QEHwE3SazDi?usp=sharing).

## Step2
Firstly,
```
cd Step2
```
Step2 here corresponds to Step2 and Step3 in the figure. Before training, please rename the model trained in Step 1 to model_epoch_-1.pth and add it to the model generation path. In addition, please download the file in this [link](https://drive.google.com/file/d/1JYqrDwsumzDsUHgwdKopVBUuAkSgILNu/view?usp=drive_link) to the Step2 folder.

### Training
For Yelp dataset,
```
python main.py --train_file ../data/yelp/yelp.train.txt \
--valid_file ../data/yelp/yelp.valid.txt \
--per_gpu_train_batch_size 4 \
--load_epoch -1 \
--cycle_annealing
```
For Yahoo dataset,
```
python main.py --train_file ../data/yahoo/yahoo.train.txt \
--valid_file ../data/yahoo/yahoo.valid.txt \
--per_gpu_train_batch_size 4 \
--load_epoch -1 \
--cycle_annealing
```

### Test
For Yelp dataset,
```
python main.py --train_file ../data/yelp/yelp.train.txt \
--valid_file ../data/yelp/yelp.valid.txt \
--per_gpu_train_batch_size 4 \
--load_epoch -1 \
--cycle_annealing \
--eval \
--eval_metrics
```
For Yahoo dataset,
```
python main.py --train_file ../data/yahoo/yahoo.train.txt \
--valid_file ../data/yahoo/yahoo.valid.txt \
--per_gpu_train_batch_size 4 \
--load_epoch -1 \
--cycle_annealing \
--eval \
--eval_metrics
```

### Generation
For Yelp dataset,
```
python main.py --generation \
--test_file ../data/yelp/yelp.test.txt \
 --load_epoch -1 \
--top_k 50 \
--top_p 0.9
```
````

---

**Machine-readable endpoints**

- JSON: [`/api/graphcanon/tools/trustedllm-regavae`](/api/graphcanon/tools/trustedllm-regavae)
- LLM index: [/llms.txt](/llms.txt)
- Full corpus: [/llms-full.txt](/llms-full.txt)

_GraphCanon - The knowledge graph for AI development. https://www.graphcanon.com/_