Improving Long Video Generation Memory

# Illustrative sketch (not from the paper)
import torch

# rotation helper: R(theta * t) @ x  (placeholder implementation)
def rotate(x, theta, t):
    # In practice R is a RoPE matrix; here we use a simple phase rotation
    angle = theta * t
    rot = torch.stack([torch.cos(angle), -torch.sin(angle), torch.sin(angle), torch.cos(angle)], dim=-1).view(x.shape[-1], x.shape[-1])
    return rot @ x

# KV cache containers
R = []                     # recent verbatim window (list of (K,V) tuples)
S = []                     # summary cache slots (list of canonical keys)
N_r = 4                    # size of recent window (example)
N_s = 3                    # number of summary slots (example)
L_attn = 5                # local attention window size (example)
mode = "field"           # "field" or "landmark"

def evict_and_compress(old_kv):
    # old_kv: list of (K, t) for the evicted chunk (M entries)
    if mode == "field":
        # Field: un‑rotate to canonical, average, re‑rotate to virtual slot position
        M = len(old_kv)
        canon_sum = torch.zeros_like(old_kv[0][0])
        for K, t in old_kv:
            canon_sum += rotate(K, -theta_f, t)   # un‑rotate
        canon_avg = canon_sum / M
        # assign to next free summary slot (simple round‑robin here)
        s = len(S) % N_s
        t_v = q - (L_attn - 1 - s)               # virtual position per paper
        S.append(rotate(canon_avg, theta_f, t_v))
    else:  # landmark mode
        # pick landmark frame (e.g., first entry) and freeze its canonical key
        K_l, t_l = old_kv[0]
        canon = rotate(K_l, -theta_f, t_l)
        s = len(S) % N_s
        t_v = q - (L_attn - 1 - s)
        S.append(rotate(canon, theta_f, t_v))

# generation step (simplified)
for q in range(10):  # q = current generation timestep
    theta_f = 0.01                     # example RoPE frequency
    # 1) produce new KV for current chunk (placeholder tensors)
    K_new = torch.randn(64)            # key vector
    V_new = torch.randn(64)            # value vector
    R.append((K_new, V_new))
    # 2) if recent window exceeds capacity, evict oldest chunk and compress
    if len(R) > N_r:
        evicted = R.pop(0)             # oldest entry (single KV for illustration)
        evict_and_compress([evicted])
    # 3) attention: rotate each summary slot to its virtual position before scoring
    attn_keys = []
    for s, K_canon in enumerate(S):
        t_v = q - (L_attn - 1 - s)
        K_rot = rotate(K_canon, theta_f, t_v)
        attn_keys.append(K_rot)
    # (attention computation omitted)

// Illustrative sketch (not from the paper)
const tf = require('@tensorflow/tfjs-node');

// rotation helper: R(theta * t) * x  (placeholder implementation)
function rotate(x, theta, t) {
  const angle = theta * t;
  const cos = Math.cos(angle);
  const sin = Math.sin(angle);
  // 2x2 rotation matrix applied to each vector slice (simplified)
  const rot = tf.tensor2d([cos, -sin, sin, cos], [2, 2]);
  return tf.matMul(rot, x);
}

// KV cache containers
let R = []; // recent verbatim window (array of {K, V})
let S = []; // summary cache slots (canonical keys)
const N_r = 4; // recent window size (example)
const N_s = 3; // number of summary slots (example)
const L_attn = 5; // local attention window size (example)
let mode = 'field'; // 'field' or 'landmark'

function evictAndCompress(oldKVs, q, theta_f) {
  if (mode === 'field') {
    // Field: un‑rotate, average, re‑rotate to virtual position
    const M = oldKVs.length;
    let canonSum = tf.zerosLike(oldKVs[0].K);
    oldKVs.forEach(({K, t}) => {
      canonSum = tf.add(canonSum, rotate(K, -theta_f, t)); // un‑rotate
    });
    const canonAvg = tf.div(canonSum, M);
    const s = S.length % N_s;
    const t_v = q - (L_attn - 1 - s);
    S.push(rotate(canonAvg, theta_f, t_v));
  } else {
    // Landmark: freeze canonical key of selected frame
    const {K, t} = oldKVs[0];
    const canon = rotate(K, -theta_f, t);
    const s = S.length % N_s;
    const t_v = q - (L_attn - 1 - s);
    S.push(rotate(canon, theta_f, t_v));
  }
}

// generation loop (simplified)
for (let q = 0; q < 10; q++) {
  const theta_f = 0.01; // example RoPE frequency
  // 1) produce new KV for current chunk (placeholder tensors)
  const K_new = tf.randomNormal([2]); // key vector (2‑dim for demo)
  const V_new = tf.randomNormal([2]); // value vector
  R.push({K: K_new, V: V_new});
  // 2) evict if recent window exceeds capacity
  if (R.length > N_r) {
    const evicted = R.shift(); // oldest entry
    evictAndCompress([evicted], q, theta_f);
  }
  // 3) attention: rotate each summary slot to its virtual position
  const attnKeys = S.map((K_canon, s) => {
    const t_v = q - (L_attn - 1 - s);
    return rotate(K_canon, theta_f, t_v);
  });
  // (attention scoring omitted)
}