A question on the discriminator detail

[johannwyh]

Hello,

I have question on the time distance information embedding in Section 3.2. Your paper says that

we encode them with positional encoding [60, 66], preprocess with a 2-layer MLP into p(δ1), ..., p(δ{k−1}) ∈ R^d and concatenate into a single vector p_δ ∈ R^(k−1)d

May I ask that what is your choice of the parameter d? And how is the positional encoding done in detail?

[universome]

Hi! We encode the positional information in 2 ways: with Fourier Features and with learnable embeddings of size 256. We concatenate the positional encoding vectors from those two modules together. You might like to wait for the complete source code release (I only need to refactor the CLIP editing scripts), but if you like, here is the positional embedding in D part:

from typing import Tuple
import torch
import torch.nn as nn
import numpy as np
from omegaconf import DictConfig

class TemporalDifferenceEncoder(nn.Module):
    def __init__(self, cfg: DictConfig):
        super().__init__()

        self.cfg = cfg

        if self.cfg.sampling.num_frames_per_video > 1:
            self.d = 256
            self.const_embed = nn.Embedding(self.cfg.sampling.max_num_frames, self.d)
            self.time_encoder = ScalarEncoder1d(self.cfg.sampling.max_num_frames)

    def get_dim(self) -> int:
        if self.cfg.sampling.num_frames_per_video == 1:
            return 1
        elif self.cfg.sampling.type == 'uniform':
            return self.d + self.time_encoder.get_dim()
        else:
            return (self.d + self.time_encoder.get_dim()) * (self.cfg.sampling.num_frames_per_video - 1)

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        batch_size = t.shape[0]

        if self.cfg.sampling.num_frames_per_video == 1:
            out = torch.zeros(len(t), 1, device=t.device)
        else:
            if self.cfg.sampling.type == 'uniform':
                num_diffs_to_use = 1
                t_diffs = t[:, 1] - t[:, 0] # [batch_size]
            else:
                num_diffs_to_use = self.cfg.sampling.num_frames_per_video - 1
                t_diffs = (t[:, 1:] - t[:, :-1]).view(-1) # [batch_size * (num_frames - 1)]
            # Note: float => round => long is necessary when it's originally long
            const_embs = self.const_embed(t_diffs.float().round().long()) # [batch_size * num_diffs_to_use, d]
            fourier_embs = self.time_encoder(t_diffs.unsqueeze(1)) # [batch_size * num_diffs_to_use, num_fourier_feats]
            out = torch.cat([const_embs, fourier_embs], dim=1) # [batch_size * num_diffs_to_use, d + num_fourier_feats]
            out = out.view(batch_size, num_diffs_to_use, -1).view(batch_size, -1) # [batch_size, num_diffs_to_use * (d + num_fourier_feats)]

        return out

#----------------------------------------------------------------------------

class ScalarEncoder1d(nn.Module):
    def __init__(self, max_num_frames: int):
        super().__init__()

        assert max_num_frames >= 1, f"Wrong max_num_frames: {max_num_frames}"
        fourier_coefs = construct_log_spaced_freqs(max_num_frames)
        self.register_buffer('fourier_coefs', fourier_coefs) # [1, num_fourier_feats]

    def get_dim(self) -> int:
        return self.fourier_coefs.shape[1] * 2

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        assert t.ndim == 2, f"Wrong shape: {t.shape}"

        t = t.view(-1).float() # [batch_size * num_frames]
        fourier_raw_embs = self.fourier_coefs * t.unsqueeze(1) # [bf, num_fourier_feats]

        fourier_embs = torch.cat([
            fourier_raw_embs.sin(),
            fourier_raw_embs.cos(),
        ], dim=1) # [bf, num_fourier_feats * 2]

        return fourier_embs

#----------------------------------------------------------------------------

def construct_log_spaced_freqs(max_num_frames: int) -> Tuple[int, torch.Tensor]:
    time_resolution = 2 ** np.ceil(np.log2(max_num_frames))
    num_fourier_feats = np.ceil(np.log2(time_resolution)).astype(int)
    powers = torch.tensor([2]).repeat(num_fourier_feats).pow(torch.arange(num_fourier_feats)) # [num_fourier_feats]
    fourier_coefs = powers.unsqueeze(0).float() * np.pi # [1, num_fourier_feats]

    return fourier_coefs / time_resolution

#----------------------------------------------------------------------------

Also note that we didn't experiment with this module much. We tried a couple of modifications, but they all performed the same.

[johannwyh]

Many thanks to your reply!

Currently I am implementing a positional encoding provided in Section 3.1 as p(t) = alpha sin(omega t + rho), with alpha, omega, rho learnable.

class PositionalEncoding(Module):
    """
    TODO: Other PE type implementation, E.g. [..., cos(2*pi*sigma^(j/m)*v), sin(2*pi*sigma^(j/m)*v), ...]
    """
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        self.omega = nn.Parameter(torch.randn(1, self.dim, requires_grad=True))
        self.rho   = nn.Parameter(torch.randn(self.dim, requires_grad=True))
        self.alpha = nn.Parameter(torch.randn(self.dim, requires_grad=True))

    def forward(self, Ts):
        """
        Ts: [b, t, 1]
        """
        ot = torch.einsum("btc,cd->btd", [Ts, self.omega])
        pe = torch.sin(ot + self.rho) * self.alpha
        return pe

Is it supposed to make sense?

[universome]

Hi! Yes, it does make sense, we tried such parametrization, and it produced very repetitive motions. Is your application also video generation?

[johannwyh]

Yes! We are following a generally different idea, but a PE of time stamps is still necessary. Many thanks for your sharing.

[universome]

Hi! We've pushed the training/evaluation/sampling code to this repo several days ago. You can check our discriminator details

[johannwyh]

Thank you for your great work! This issue can be closed.