### Complete GAN Training Loop Example

Source: https://context7.com/lucidrains/transformer-lm-gan/llms.txt

Demonstrates a full training loop for the Transformer LM GAN. It includes discriminator training with hinge loss and generator training via an evolution strategy for gradient-free optimization. Requires custom dataset and optimizer setup.

```python
import torch
from torch.optim import Adam
from torch.utils.data import DataLoader, Dataset
from x_evolution import EvoStrategy
from transformer_lm_gan import LanguageModelGenerator, Discriminator
from transformer_lm_gan.transformer_lm_gan import discriminator_hinge_loss

# Configuration
BATCH_SIZE = 4
SEQ_LEN = 64
LEARNING_RATE = 1e-4
NUM_ITERATIONS = 1000
GRAD_ACCUM_EVERY = 2

# Initialize models
generator = LanguageModelGenerator(
    num_tokens=256, dim=512, depth=6, dim_head=64, heads=8, max_seq_len=SEQ_LEN
).cuda()

discriminator = Discriminator(
    num_tokens=256, dim=512, depth=2, dim_head=64, heads=9, max_seq_len=SEQ_LEN
).cuda()

# Custom dataset (example with random data - replace with real text data)
class TextDataset(Dataset):
    def __init__(self, size=10000, seq_len=64):
        self.data = torch.randint(0, 256, (size, seq_len + 1))
    def __len__(self):
        return len(self.data)
    def __getitem__(self, idx):
        return self.data[idx].long()

dataset = TextDataset(size=10000, seq_len=SEQ_LEN)
dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)

# Fitness function for evolution strategy
def get_fitness(model):
    model.eval()
    rand_token = torch.randint(0, 256, (BATCH_SIZE, 1), device='cuda')
    generated, _ = model.model.generate(rand_token, SEQ_LEN, disable_pbar=True)
    discriminator.eval()
    with torch.no_grad():
        score = discriminator(generated)
    return -score.mean()  # Negative because evolution maximizes

# Evolution strategy for generator
evo_strat = EvoStrategy(
    generator,
    environment=get_fitness,
    noise_population_size=5,
    num_generations=NUM_ITERATIONS,
    learning_rate=LEARNING_RATE,
    use_optimizer=True,
    verbose=False
)

# Discriminator optimizer
discr_optim = Adam(discriminator.parameters(), lr=LEARNING_RATE)
```

--------------------------------

### Initialize and Use LanguageModelGenerator

Source: https://context7.com/lucidrains/transformer-lm-gan/llms.txt

Demonstrates initializing the generator, computing autoregressive loss, and performing inference with KV-caching.

```python
import torch
from transformer_lm_gan import LanguageModelGenerator

# Initialize the generator with transformer architecture parameters
generator = LanguageModelGenerator(
    num_tokens=256,      # Vocabulary size (e.g., 256 for byte-level)
    dim=512,             # Model dimension
    depth=6,             # Number of transformer layers
    dim_head=64,         # Dimension per attention head
    heads=8,             # Number of attention heads
    max_seq_len=1024     # Maximum sequence length
)

generator = generator.cuda()

# Forward pass for computing autoregressive loss
input_seq = torch.randint(0, 256, (2, 512)).cuda()  # (batch, seq_len)
ar_loss = generator(input_seq, return_ar_loss=True)
print(f"Autoregressive loss: {ar_loss.item():.4f}")

# Forward pass for getting logits
logits = generator(input_seq)  # Returns (batch, seq_len, num_tokens)
print(f"Logits shape: {logits.shape}")

# Forward pass with cached key-values for efficient inference
logits, cache = generator(input_seq, return_intermediates=True)
# Use cache for subsequent tokens to avoid recomputation
next_token = torch.randint(0, 256, (2, 1)).cuda()
next_logits, updated_cache = generator(next_token, return_intermediates=True, cache=cache)
```

--------------------------------

### Initialize and Train GAN for Language Modeling

Source: https://github.com/lucidrains/transformer-lm-gan/blob/main/README.md

Configures the GAN with generator and discriminator parameters, then performs forward passes for both components on a sequence.

```python
import torch

from transformer_lm_gan import (
    LanguageModelGenerator,
    Discriminator,
    GAN,
)

gan = GAN(
    strategy = 'gumbel_one_hot', # or 'rotate' for rotation trick, may try combination of two if both fails in experiments
    generator = dict(
        num_tokens = 256,
        dim = 512,
        depth = 6,
        dim_head = 64,
        heads = 8,
        max_seq_len = 1024
    ),
    discriminator = dict(
        num_tokens = 256,
        dim = 512,
        depth = 2,
        dim_head = 64,
        heads = 9,
        max_seq_len = 1024
    )
).cuda()

seq = torch.randint(0, 256, (2, 1024)).cuda()

discr_loss = gan.discriminate_forward(seq)
discr_loss.backward()

gen_loss = gan.generate_forward(seq)
gen_loss.backward()
```

--------------------------------

### Initialize Discriminator

Source: https://context7.com/lucidrains/transformer-lm-gan/llms.txt

Basic import for the Discriminator class.

```python
import torch
from transformer_lm_gan import Discriminator
```

--------------------------------

### Generate Text with LanguageModelGenerator

Source: https://context7.com/lucidrains/transformer-lm-gan/llms.txt

Shows how to use the generate method with sampling parameters and KV-caching for autoregressive text production.

```python
import torch
from transformer_lm_gan import LanguageModelGenerator

generator = LanguageModelGenerator(
    num_tokens=256,
    dim=512,
    depth=6,
    dim_head=64,
    heads=8,
    max_seq_len=1024
).cuda()

# Create a prompt (batch_size=1, prompt_length=10)
prompt = torch.randint(0, 256, (1, 10)).cuda()

# Generate text with default parameters
generated_seq, sampling_params = generator.generate(
    prompt=prompt,
    seq_len=100,           # Total sequence length (including prompt)
    temperature=1.0,       # Sampling temperature (higher = more random)
    filter_thres=0.9,      # Top-k threshold (keep top 10% of tokens)
    cache_kv=True,         # Use KV-cache for faster generation
    return_with_prompt=True,  # Include prompt in output
    disable_pbar=False     # Show progress bar
)

print(f"Generated sequence shape: {generated_seq.shape}")  # (1, 100)
print(f"Sampling params: filter_fn, filter_thres={sampling_params[1]}, temp={sampling_params[2]}")

# Generate without prompt in output
generated_only, _ = generator.generate(
    prompt=prompt,
    seq_len=50,
    return_with_prompt=False
)
print(f"Generated only shape: {generated_only.shape}")  # (1, 40) = 50 - 10

# Decode byte-level tokens to text
def decode_tokens(tokens):
    return "".join([chr(max(32, t)) for t in tokens.tolist()])

decoded_text = decode_tokens(generated_seq[0])
print(f"Generated text: {decoded_text}")
```

--------------------------------

### Initialize and Use Discriminator

Source: https://context7.com/lucidrains/transformer-lm-gan/llms.txt

Initializes the Discriminator model and demonstrates forward passes with both discrete token indices and continuous embeddings. The discriminator outputs logits indicating the perceived 'realness' of the input sequences.

```python
discriminator = Discriminator(
    num_tokens=256,      # Vocabulary size
    dim=512,             # Model dimension
    depth=2,             # Number of transformer layers (typically smaller than generator)
    dim_head=64,         # Dimension per attention head
    heads=9,             # Number of attention heads
    max_seq_len=1024,    # Maximum sequence length
    gp_weight=10.0       # Gradient penalty weight (for WGAN-GP style training)
)

discriminator = discriminator.cuda()

# Forward pass with discrete token indices
real_sequences = torch.randint(0, 256, (4, 128)).cuda()  # (batch, seq_len)
real_logits = discriminator(real_sequences)
print(f"Real logits shape: {real_logits.shape}")  # (4,)
print(f"Real logits: {real_logits}")

# Forward pass with continuous embeddings (for gradient flow from generator)
continuous_embeddings = torch.randn(4, 128, 512).cuda()  # (batch, seq_len, dim)
fake_logits = discriminator(continuous_embeddings)
print(f"Fake logits: {fake_logits}")

# Higher logits indicate "more real"
print(f"Discriminator predicts real: {(real_logits > 0).tolist()}")
```

--------------------------------

### Reference Citations

Source: https://github.com/lucidrains/transformer-lm-gan/blob/main/README.md

BibTeX citations for the research papers related to the implementation.

```bibtex
@inproceedings{Huang2025TheGI,
    title   = {The GAN is dead; long live the GAN! A Modern GAN Baseline},
    author  = {Yiwen Huang and Aaron Gokaslan and Volodymyr Kuleshov and James Tompkin},
    year    = {2025},
    url     = {https://api.semanticscholar.org/CorpusID:275405495}
}
```

```bibtex
@misc{sarkar2025evolutionstrategieshyperscale,
    title     = {Evolution Strategies at the Hyperscale}, 
    author    = {Bidipta Sarkar and Mattie Fellows and Juan Agustin Duque and Alistair Letcher and Antonio León Villares and Anya Sims and Dylan Cope and Jarek Liesen and Lukas Seier and Theo Wolf and Uljad Berdica and Alexander David Goldie and Aaron Courville and Karin Sevegnani and Shimon Whiteson and Jakob Nicolaus Foerster},
    year      = {2025},
    eprint    = {2511.16652},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG},
    url       = {https://arxiv.org/abs/2511.16652}, 
}
```

--------------------------------

### Adversarial Training Loop

Source: https://context7.com/lucidrains/transformer-lm-gan/llms.txt

Executes the training cycle for both the discriminator and the generator using gradient accumulation and evolution strategies.

```python
data_iter = iter(dataloader)
for i in range(NUM_ITERATIONS):
    # Train discriminator
    discriminator.train()
    generator.eval()

    for _ in range(GRAD_ACCUM_EVERY):
        try:
            real_data = next(data_iter).cuda()
        except StopIteration:
            data_iter = iter(dataloader)
            real_data = next(data_iter).cuda()

        # Generate fake data
        prompts = real_data[:, :1]
        with torch.no_grad():
            fake_data, _ = generator.generate(prompts, SEQ_LEN, disable_pbar=True)

        # Discriminator loss
        real_logits = discriminator(real_data)
        fake_logits = discriminator(fake_data)
        loss = discriminator_hinge_loss(real_logits, fake_logits)
        (loss / GRAD_ACCUM_EVERY).backward()

    torch.nn.utils.clip_grad_norm_(discriminator.parameters(), 0.5)
    discr_optim.step()
    discr_optim.zero_grad()

    # Train generator with evolution strategy
    evo_strat(num_generations=1)

    if i % 100 == 0:
        print(f"Iter {i}: discriminator loss = {loss.item():.4f}")
```

--------------------------------

### LanguageModelGenerator Initialization and Forward Pass

Source: https://context7.com/lucidrains/transformer-lm-gan/llms.txt

Initializes the transformer-based generator and performs forward passes for autoregressive loss calculation or logit generation.

```APIDOC
## LanguageModelGenerator Initialization

### Description
Initializes a causal transformer decoder for autoregressive text generation.

### Parameters
- **num_tokens** (int) - Required - Vocabulary size.
- **dim** (int) - Required - Model dimension.
- **depth** (int) - Required - Number of transformer layers.
- **dim_head** (int) - Required - Dimension per attention head.
- **heads** (int) - Required - Number of attention heads.
- **max_seq_len** (int) - Required - Maximum sequence length.

### Request Example
```python
generator = LanguageModelGenerator(num_tokens=256, dim=512, depth=6, dim_head=64, heads=8, max_seq_len=1024)
```
```

--------------------------------

### LanguageModelGenerator.generate

Source: https://context7.com/lucidrains/transformer-lm-gan/llms.txt

Performs autoregressive text generation from a prompt sequence with support for sampling strategies and KV-caching.

```APIDOC
## generate

### Description
Generates text sequences token by token based on a provided prompt.

### Parameters
- **prompt** (torch.Tensor) - Required - Input sequence tensor.
- **seq_len** (int) - Required - Total sequence length including prompt.
- **temperature** (float) - Optional - Sampling temperature.
- **filter_thres** (float) - Optional - Top-k filtering threshold.
- **cache_kv** (bool) - Optional - Whether to use KV-caching.
- **return_with_prompt** (bool) - Optional - Whether to include the prompt in the output.

### Response
- **generated_seq** (torch.Tensor) - The generated sequence.
- **sampling_params** (tuple) - Parameters used during sampling.
```

--------------------------------

### Compute Discriminator Hinge Loss

Source: https://context7.com/lucidrains/transformer-lm-gan/llms.txt

Calculates the discriminator's hinge loss using real and fake data logits. This loss function encourages the discriminator to output values greater than 1 for real data and less than -1 for fake data, ensuring stable gradients.

```python
import torch
import torch.nn.functional as F
from transformer_lm_gan import LanguageModelGenerator, Discriminator
from transformer_lm_gan.transformer_lm_gan import discriminator_hinge_loss

# Initialize models
generator = LanguageModelGenerator(
    num_tokens=256, dim=512, depth=6, dim_head=64, heads=8, max_seq_len=128
).cuda()

discriminator = Discriminator(
    num_tokens=256, dim=512, depth=2, dim_head=64, heads=9, max_seq_len=128
).cuda()

# Prepare real data
real_data = torch.randint(0, 256, (4, 128)).cuda()

# Generate fake data
prompts = real_data[:, :1]  # Use first token as prompt
with torch.no_grad():
    fake_data, _ = generator.generate(prompts, seq_len=128, disable_pbar=True)

# Compute discriminator outputs
real_logits = discriminator(real_data)
fake_logits = discriminator(fake_data)

# Compute hinge loss: wants real > 1 and fake < -1
loss = discriminator_hinge_loss(real_logits, fake_logits)
print(f"Discriminator hinge loss: {loss.item():.4f}")

# The hinge loss formula: mean(relu(1 + real) + relu(1 - fake))
# Optimal when real_logits > 1 and fake_logits < -1
manual_loss = (F.relu(1 + real_logits) + F.relu(1 - fake_logits)).mean()
print(f"Manual computation: {manual_loss.item():.4f}")

# Backpropagate
loss.backward()
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.