Orthogonal seq gen.

gistfile1.txt

def random_so(batch, n, *, generator, device):
    matrix = torch.randn((batch, n, n), generator=generator, device=device)
    q, r = torch.linalg.qr(matrix)

    # Sign-correct QR to obtain Haar-random orthogonal matrices.
    diagonal = torch.diagonal(r, dim1=-2, dim2=-1)
    signs = torch.where(diagonal >= 0, 1.0, -1.0)
    q = q * signs.unsqueeze(-2)

    # Force determinant +1, giving SO(N) rather than O(N).
    determinants = torch.linalg.det(q)
    q[:, :, 0] *= torch.where(
        determinants < 0, -1.0, 1.0
    ).unsqueeze(-1)
    return q


@torch.inference_mode()
def generate_batch(lengths, *, n, vocab_size, generator, device):
    batch = len(lengths)
    max_length = int(lengths.max())

    # Independently sampled for every document.
    T = random_so(batch, n, generator=generator, device=device)

    h = torch.randn((batch, n), generator=generator, device=device)
    h = h / torch.linalg.vector_norm(h, dim=-1, keepdim=True)

    R = (
        torch.randn(
            (batch, vocab_size, n),
            generator=generator,
            device=device,
        )
        / math.sqrt(n)
    )

    output = torch.empty(
        (batch, max_length),
        dtype=torch.uint8,
        device=device,
    )

    for position in range(max_length):
        logits = torch.bmm(R, h.unsqueeze(-1)).squeeze(-1)
        output[:, position] = logits.argmax(dim=-1).to(torch.uint8)
        h = torch.bmm(T, h.unsqueeze(-1)).squeeze(-1)

    return output.cpu().numpy()