I'm hoping to use quimb for a quick diagnostic. However, I have been unable to use MERA with a torch backend so this can run on GPU. I've explicitly set backend as much as possible, and I explicitly move MPS and MERA tensors to torch during initialization. When I contract the 2, however, a new TensorNetwork object is made, and the tensors are converted to numpy somewhere along the way. Is there a way to prevent this from happening?
My code, with a TODO pointing to where things break
```python
"""MERA compression using quimb tensor network library.
Correct approach: Treat data [seq_len, hidden_dim] as a product-state MPS
where each site has physical dimension = hidden_dim.
NO FLATTENING - preserves structure and avoids exponential blowup.
"""
import torch
import autoray
import quimb.tensor as qtn
autoray.set_backend("torch")
qtn.set_contract_backend("torch")
qtn.contraction.set_contract_backend("torch")
def train_mera_on_data(mera, data, num_steps=100, optimizer="adam", backend="torch"):
"""Train MERA to compress activation data.
Args:
mera: Initialized MERA tensor network
data: Tensor [batch, seq_len, hidden_dim]
num_steps: Optimization steps
optimizer: Optimizer name
backend: Autodiff backend
Returns:
Optimized MERA tensor network
"""
batch_size, seq_len, hidden_dim = data.shape
print(f"=" * 70)
print(f"MERA TRAINING")
print(f"=" * 70)
print(f" Data: [{batch_size}, {seq_len}, {hidden_dim}]")
print(f" Optimizer: {optimizer}")
print(f" Steps: {num_steps}")
# Extract single sample and normalize
target_matrix = data[0].detach().double()
target_matrix = target_matrix / torch.linalg.norm(target_matrix)
print(f"\nTarget data:")
print(f" Shape: {target_matrix.shape}")
print(f" Norm: {torch.linalg.norm(target_matrix).item():.6f}")
print(f" Device: {target_matrix.device}")
# Convert to tensor network
print(f"\nCreating data tensor network...")
data_tn = data_to_tensor_network(target_matrix, use_torch=True)
print(f" Number of tensors: {len(data_tn.tensor_map)}")
# Loss function: tensor network contraction for overlap
def loss_fn(mera_tn):
"""Compute overlap via tensor network contraction.
Both MERA and data TN share physical indices k0, k1, ..., k{L-1}.
Contracting these gives the inner product ⟨MERA|data⟩.
"""
# Contract MERA with data TN
# TODO: numpy / torch collision at this point,
# even though mera_tn and data_tn
overlap = mera_tn @ data_tn
# # Contract entire network to scalar - explicitly use torch backend
# overlap = inner_product_tn.contract(all, optimize="greedy", backend="torch")
# Loss: 1 - |overlap|²
loss = 1.0 - abs(overlap) ** 2
return loss
# Norm function: skip normalization to avoid expensive contraction
# (normalization can be done periodically instead of every step)
def norm_fn(mera_tn):
return mera_tn # Identity - no normalization
# Create optimizer - TNOptimizer will extract variables automatically
print(f"\nSetting up TNOptimizer...")
tnopt = qtn.TNOptimizer(
mera,
loss_fn=loss_fn,
norm_fn=norm_fn,
optimizer=optimizer,
autodiff_backend=backend,
)
print(f" Number of parameters: {sum(t.size for t in mera.tensor_map.values())}")
# Run optimization
print(f"\nOptimizing for {num_steps} steps...")
initial_loss = loss_fn(mera)
print(f" Initial loss: {initial_loss:.6f}")
optimized_params = tnopt.optimize(n=num_steps)
# Get optimized MERA - it's stored in tnopt.tn
optimized_mera = tnopt.get_tn_opt()
final_loss = loss_fn(optimized_mera)
print(f"\nOptimization complete!")
print(f" Final loss: {final_loss:.6f}")
print(f" Loss reduction: {initial_loss - final_loss:.6f}")
print(f" Final overlap: {torch.sqrt(1.0 - final_loss):.6f}")
# Compute compression ratio
original_size = seq_len * hidden_dim
mera_params = sum(t.size for t in optimized_mera.tensors)
compression_ratio = original_size / mera_params
print(f"\nCompression analysis:")
print(f" Original parameters: {original_size}")
print(f" MERA parameters: {mera_params}")
print(f" Compression ratio: {compression_ratio:.2f}x")
return optimized_mera
def generate_power_law_tree_dataset(
batch_size: int = 32,
seq_len: int = 4096,
hidden_dim: int = 128,
alpha: float = 0.5,
device: str = "cpu",
) -> torch.Tensor:
"""
Generates a synthetic dataset of shape [batch_size, seq_len, hidden_dim]
where the correlations along the sequence dimension follow a power law
decay dictated by alpha.
This matches a Tree Tensor Network / MERA generative model.
"""
assert (
seq_len & (seq_len - 1)
) == 0, "seq_len must be a power of 2 for hierarchical generation."
num_layers = int(torch.log2(torch.tensor([seq_len])).item())
# Start with a global root state for each sample and channel
# Shape: [batch_size, 1, hidden_dim]
current_state = torch.randn(batch_size, 1, hidden_dim, device=device)
# Set up random orthogonal projection matrices to simulate local mixing (isometries)
# This prevents the channels from being trivially decoupled
proj = torch.randn(hidden_dim, hidden_dim, device=device)
Q, _ = torch.linalg.qr(proj)
# Hierarchically split and inject scale-dependent power-law variance
for layer in range(num_layers):
# Power-law variance scaling factor for this specific length-scale
# As layer increases (moving to finer localized scales), noise amplitude shrinks
scale_factor = 1.0 / ((layer + 1) ** alpha)
# Duplicate the sequence length by splitting each node into a pair
# [batch_size, current_len, hidden_dim] -> [batch_size, current_len, 2, hidden_dim]
current_state = current_state.unsqueeze(2).repeat(1, 1, 2, 1)
# Inject fine-grained details at this scale
noise = torch.randn_like(current_state) * scale_factor
current_state = current_state + noise
# Collapse back to a flat sequence for the next tree layer
batch, curr_len, pairs, h_dim = current_state.shape
current_state = current_state.view(batch, curr_len * pairs, h_dim)
# Apply a mixing isometry across channels to blend the features linearly
current_state = current_state @ Q
# Final pass: Center and normalize to prevent variance explosion
current_state = current_state - current_state.mean(dim=1, keepdim=True)
current_state = current_state / (current_state.std(dim=1, keepdim=True) + 1e-8)
return current_state
def data_to_tensor_network(data_matrix, use_torch=True):
"""Convert data [seq_len, hidden_dim] to tensor network (product-state MPS).
Each site i gets a tensor with physical index k{i} of dimension hidden_dim.
Bond dimension between sites = 1 (product state, no entanglement).
Args:
data_matrix: Array/Tensor [seq_len, hidden_dim]
use_torch: If True, ensure tensors are torch tensors
Returns:
TensorNetwork representing the data
"""
# Convert to torch if needed
if use_torch and not isinstance(data_matrix, torch.Tensor):
data_matrix = torch.from_numpy(data_matrix)
seq_len, hidden_dim = data_matrix.shape
data_tensors = []
for i in range(seq_len):
# Local token vector of shape (hidden_dim,)
vec = data_matrix[i]
# Ensure it's a torch tensor
if use_torch and not isinstance(vec, torch.Tensor):
vec = torch.from_numpy(vec)
# Create tensor for this site with physical index k{i}
t = qtn.Tensor(data=vec, inds=(f"k{i}",), tags={f"I{i}", "DATA"})
data_tensors.append(t)
# Combine into unified tensor network
data_tn = qtn.TensorNetwork(data_tensors)
return data_tn
def initialize_mera(seq_len, hidden_dim, max_bond, device):
"""Initialize MERA tensor network.
Args:
seq_len: Sequence length (must be power of 2)
hidden_dim: Physical dimension per site
max_bond: Maximum bond dimension
device: torch device
Returns:
MERA tensor network with torch tensors
"""
assert (seq_len & (seq_len - 1)) == 0, "seq_len must be power of 2"
print(f"\nInitializing MERA...")
print(f" Sites: {seq_len}")
print(f" Physical dim: {hidden_dim}")
print(f" Max bond: {max_bond}")
mera = qtn.MERA.rand(
L=seq_len, phys_dim=hidden_dim, D=max_bond, dtype="float64", backend="torch"
)
# Ensure all tensors are torch on correct device
for tid, tensor in mera.tensor_map.items():
if isinstance(tensor.data, torch.Tensor):
tensor.modify(data=tensor.data.to(device))
else:
tensor.modify(data=torch.from_numpy(tensor.data).to(device))
mera.tensor_map[tid] = tensor
print(f" Number of tensors: {len(mera.tensor_map)}")
print(f" Types: {[type(a) for a in mera.arrays]}")
print(f" Device: {device}")
return mera
if name == "main":
# Generate tree-structured test data
print("\nGenerating tree-structured dataset (α=1.0)...")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
data = generate_power_law_tree_dataset(
batch_size=2,
seq_len=8, # Very small to reduce memory
hidden_dim=4, # Very small dimension
alpha=1.0,
device=device,
)
print(f"Generated data shape: {data.shape}")
# Initialize MERA
mera = initialize_mera(
seq_len=data.shape[1], hidden_dim=data.shape[2], max_bond=4, device=device
)
# Train MERA
mera_opt = train_mera_on_data(
mera=mera,
data=data,
num_steps=10,
optimizer="adam",
backend="torch",
)
print("\n" + "=" * 70)
print("DONE")
print("=" * 70)
What is your issue?
I would like to see if activations pulled from a large language model (shape
[batch_size, sequence_length, hidden_dim]) are compressible with a MERA tensor network. The hypothesis is that the power-law(ish) correlation along the sequence_length dim is similar what would happen to correlation length at a quantum critical point, so it might be compressible with the MERA ansatz.I'm hoping to use quimb for a quick diagnostic. However, I have been unable to use MERA with a torch backend so this can run on GPU. I've explicitly set backend as much as possible, and I explicitly move MPS and MERA tensors to torch during initialization. When I contract the 2, however, a new TensorNetwork object is made, and the tensors are converted to numpy somewhere along the way. Is there a way to prevent this from happening?
Stack Trace
My code, with a TODO pointing to where things break
```python """MERA compression using quimb tensor network library.Correct approach: Treat data [seq_len, hidden_dim] as a product-state MPS
where each site has physical dimension = hidden_dim.
NO FLATTENING - preserves structure and avoids exponential blowup.
"""
import torch
import autoray
import quimb.tensor as qtn
autoray.set_backend("torch")
qtn.set_contract_backend("torch")
qtn.contraction.set_contract_backend("torch")
def train_mera_on_data(mera, data, num_steps=100, optimizer="adam", backend="torch"):
"""Train MERA to compress activation data.
def generate_power_law_tree_dataset(
batch_size: int = 32,
seq_len: int = 4096,
hidden_dim: int = 128,
alpha: float = 0.5,
device: str = "cpu",
) -> torch.Tensor:
"""
Generates a synthetic dataset of shape [batch_size, seq_len, hidden_dim]
where the correlations along the sequence dimension follow a power law
decay dictated by alpha.
def data_to_tensor_network(data_matrix, use_torch=True):
"""Convert data [seq_len, hidden_dim] to tensor network (product-state MPS).
def initialize_mera(seq_len, hidden_dim, max_bond, device):
"""Initialize MERA tensor network.
if name == "main":
# Generate tree-structured test data
print("\nGenerating tree-structured dataset (α=1.0)...")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")