### Install UltraMem

Source: https://github.com/lucidrains/ultra-mem/blob/main/README.md

Install the UltraMem library using pip.

```shell
pip install ultra-mem
```

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

### Define Transformer with UltraMem Integration

Source: https://context7.com/lucidrains/ultra-mem/llms.txt

Defines a transformer model that incorporates UltraMem between decoder layers. This setup is suitable for language modeling tasks where memory augmentation is desired. Ensure UltraMem is initialized with the total number of transformer layers for proper variance initialization.

```python
import torch
from torch import nn
from torch.nn import Module, RMSNorm
import torch.nn.functional as F
from ultra_mem import UltraMem

class TransformerWithUltraMem(Module):
    def __init__(
        self,
        num_tokens: int,
        dim: int,
        depth: tuple[int, int, int],  # (pre, mid, post) decoder depths
        num_memories: int = 1_000_000,
    ):
        super().__init__()
        self.token_emb = nn.Embedding(num_tokens, dim)

        pre_depth, mid_depth, post_depth = depth
        total_layers = pre_depth + mid_depth + post_depth

        # UltraMem positioned between decoder stages
        self.ultra_mem = UltraMem(
            dim=dim,
            layers_for_mem_init=total_layers,  # For proper variance initialization
            num_memories=num_memories,
            core_heads=2,
            topk=32,
        )

        # Simple transformer layers (replace with actual decoder blocks)
        self.pre_layers = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(dim, nhead=8, batch_first=True),
            num_layers=pre_depth
        )
        self.mid_layers = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(dim, nhead=8, batch_first=True),
            num_layers=mid_depth
        )
        self.post_layers = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(dim, nhead=8, batch_first=True),
            num_layers=post_depth
        )

        self.norm = RMSNorm(dim)
        self.to_logits = nn.Linear(dim, num_tokens, bias=False)

    def forward(self, x, return_loss=False):
        if return_loss:
            x, labels = x[:, :-1], x[:, 1:]

        x = self.token_emb(x)

        # Pre-memory decoder layers
        x = self.pre_layers(x)

        # Memory lookup and mid decoder
        mem_out, mem_indices, aux_loss = self.ultra_mem(x)
        x = self.mid_layers(x)
        x = x + mem_out  # Residual connection with memory output

        # Post-memory decoder layers
        x = self.post_layers(x)

        logits = self.to_logits(self.norm(x))

        if not return_loss:
            return logits

        # Combine cross-entropy loss with auxiliary loss
        ce_loss = F.cross_entropy(logits.transpose(1, 2), labels)
        total_loss = ce_loss + aux_loss

        return total_loss, (ce_loss, aux_loss)


# Create model for character-level language modeling
model = TransformerWithUltraMem(
    num_tokens=256,           # Byte-level vocabulary
    dim=512,
    depth=(3, 2, 3),          # 8 total layers with UltraMem after layer 3
    num_memories=1_000_000,
)

# Training forward pass
tokens = torch.randint(0, 256, (4, 513))  # (batch, seq_len+1)
total_loss, (ce_loss, aux_loss) = model(tokens, return_loss=True)

print(f"Total loss: {total_loss.item():.4f}")
print(f"CE loss: {ce_loss.item():.4f}")
print(f"Aux loss: {aux_loss.item():.4f}")

# Inference forward pass
tokens = torch.randint(0, 256, (1, 512))
logits = model(tokens)
print(f"Logits shape: {logits.shape}")  # torch.Size([1, 512, 256])
```

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

### Initialize and Use UltraMem

Source: https://github.com/lucidrains/ultra-mem/blob/main/README.md

Initialize the UltraMem model and process input tokens. Ensure torch is imported.

```python
import torch
from ultra_mem import UltraMem

ultra_mem = UltraMem(
    dim = 512,
    core_heads = 2,
    topk = 32,
)

tokens = torch.randn(1, 1024, 512)

out, mem_indices, aux_loss = ultra_mem(tokens) # (1, 1024, 512), (2, 1, 1024, 32), ()

```

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

### Initialize and Use UltraMem Module

Source: https://context7.com/lucidrains/ultra-mem/llms.txt

Basic initialization of the UltraMem module with default settings for a million memory slots. The forward pass returns output, memory indices, and an auxiliary loss for training stability.

```python
import torch
from ultra_mem import UltraMem

# Basic initialization with 1 million memory slots
ultra_mem = UltraMem(
    dim=512,                    # Input/output dimension
    core_heads=2,               # Number of memory heads
    topk=32,                    # Top-k memories to retrieve per head
)

# Forward pass with random token embeddings
tokens = torch.randn(1, 1024, 512)  # (batch, seq_len, dim)
out, mem_indices, aux_loss = ultra_mem(tokens)

# Output shapes:
# out: (1, 1024, 512) - same shape as input
# mem_indices: (2, 1, 1024, 32) - (heads, batch, seq, topk)
# aux_loss: scalar tensor for training stability

```

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

### Sparse Finetuning with UltraMem

Source: https://context7.com/lucidrains/ultra-mem/llms.txt

Configure UltraMem for sparse finetuning by providing a trainable mask. Only memories indicated by the mask will receive gradients during training, enabling efficient adaptation.

```python
import torch
from ultra_mem import UltraMem

ultra_mem = UltraMem(
    dim=512,
    core_heads=2,
    topk=32,
    num_memories=1_000_000,
)

# Create a sparse mask indicating which memories are trainable
# Shape: (heads, num_virtual_memories)
# num_virtual_memories = num_memories // value_expansion
trainable_mask = torch.randint(0, 2, (2, ultra_mem.num_virtual_mems)).bool()

print(f"Total virtual memories: {ultra_mem.num_virtual_mems}")
print(f"Trainable memories per head: {trainable_mask.sum(dim=-1).tolist()}")

tokens = torch.randn(1, 1024, 512)
out, mem_indices, aux_loss = ultra_mem(
    tokens,
    trainable_sparse_mask=trainable_mask  # Only masked memories get gradients
)


```

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

### Advanced UltraMem Configuration

Source: https://context7.com/lucidrains/ultra-mem/llms.txt

Initialize UltraMem with extensive configuration options for custom memory behavior, including dimensions, Tucker decomposition rank, auxiliary loss settings, and gating mechanisms. The forward pass demonstrates output shapes and auxiliary loss.

```python
import torch
from torch import nn
from ultra_mem import UltraMem

# Advanced configuration with all options
ultra_mem = UltraMem(
    dim=512,                        # Input dimension
    dim_out=512,                    # Output dimension (defaults to dim)
    dim_values=256,                 # Dimension of memory values
    num_memories=1_000_000,         # Total memory slots (must be perfect square)
    topk=32,                        # Top-k memories retrieved per head
    dim_queries_keys=128,           # Query/key projection dimension
    core_rank=2,                    # Tucker decomposition rank
    core_heads=2,                   # Number of parallel memory heads
    core_aux_loss_margin=0.15,      # Margin for auxiliary loss
    aux_loss_weight=0.1,            # Weight for auxiliary loss
    score_activation=nn.ReLU(),     # Non-competitive scoring (Csordas et al.)
    value_expansion=4,              # Expand effective memory capacity 4x
    pre_query_causal_conv=True,     # Apply causal convolution before queries
    query_conv_kernel_size=5,       # Kernel size for causal convolution
    qk_layernorm=True,              # Apply LayerNorm to queries and keys
    prenorm=True,                   # Apply RMSNorm to input
    gate_values_with_input=True,    # SiLU gating on memory output
    layers_for_mem_init=12,         # Total layers (for memory initialization variance)
    mem_lr_scale=10.0,              # Memory learning rate multiplier
    mem_decay_lr_over_steps=20_000, # Steps to decay memory LR from 10x to 1x
)

tokens = torch.randn(2, 512, 512)  # (batch=2, seq_len=512, dim=512)
out, mem_indices, aux_loss = ultra_mem(tokens)

print(f"Output shape: {out.shape}")           # torch.Size([2, 512, 512])
print(f"Indices shape: {mem_indices.shape}")  # torch.Size([2, 2, 512, 32])
print(f"Auxiliary loss: {aux_loss.item():.4f}")

```

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

### UltraMem Training Loop with Auxiliary Loss

Source: https://context7.com/lucidrains/ultra-mem/llms.txt

Demonstrates a training loop for UltraMem, including gradient accumulation and the use of an auxiliary loss to maintain Tucker decomposition quality. Key parameters like `aux_loss_weight`, `mem_lr_scale`, and `mem_decay_lr_over_steps` are configurable. The `reset_step_()` method should be called if the training step counter needs to be reset.

```python
import torch
from torch.optim import Adam
from ultra_mem import UltraMem

# Training hyperparameters
LEARNING_RATE = 1e-4
GRAD_ACCUM_STEPS = 4

# Create UltraMem module
ultra_mem = UltraMem(
    dim=512,
    core_heads=2,
    topk=32,
    num_memories=1_000_000,
    aux_loss_weight=0.1,        # Weight for auxiliary loss
    mem_lr_scale=10.0,          # Memories need higher learning rate
    mem_decay_lr_over_steps=20_000,
)

optimizer = Adam(ultra_mem.parameters(), lr=LEARNING_RATE)

# Simulated training loop
ultra_mem.train()
for step in range(100):
    optimizer.zero_grad()

    # Gradient accumulation
    total_aux_loss = 0
    for _ in range(GRAD_ACCUM_STEPS):
        tokens = torch.randn(4, 512, 512)  # Simulated embeddings

        out, mem_indices, aux_loss = ultra_mem(tokens)

        # In real training, combine with main task loss
        # loss = task_loss + aux_loss
        (aux_loss / GRAD_ACCUM_STEPS).backward()
        total_aux_loss += aux_loss.item()

    torch.nn.utils.clip_grad_norm_(ultra_mem.parameters(), 0.5)
    optimizer.step()

    if step % 10 == 0:
        print(f"Step {step}, Aux loss: {total_aux_loss:.4f}, "
              f"Mem LR scale: {ultra_mem.mem_lr_scale:.2f}")

# Reset training step counter if needed
ultra_mem.reset_step_()
```

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

### Train Character-level LM

Source: https://github.com/lucidrains/ultra-mem/blob/main/README.md

Run the training script for a character-level language model.

```shell
uv run train.py
```

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