Embeddings¶

Declarative embedding pipelines for observation processing.

Overview¶

The falcon.embeddings package provides tools for building observation embedding networks via declarative YAML configuration. Embeddings map raw observations to lower-dimensional summary statistics before they enter the estimator.

Declarative Configuration¶

Embedding networks are defined in YAML using the _target_ / _input_ system:

_target_: Import path of the nn.Module class to instantiate
_input_: List of observation node names (or nested sub-embeddings) that feed into this module

Basic Embedding¶

embedding:
  _target_: model.MyEmbedding
  _input_: [x]

Multi-Input Embedding¶

embedding:
  _target_: model.MyEmbedding
  _input_: [x, y]  # Multiple observation nodes

Nested Pipeline¶

Sub-embeddings can be nested arbitrarily deep. Each _input_ entry can itself be a _target_ / _input_ block:

embedding:
  _target_: model.Concatenate
  _input_:
    - _target_: timm.create_model
      model_name: resnet18
      pretrained: true
      num_classes: 0
      _input_:
        _target_: model.Unsqueeze
        _input_: [image]
    - _target_: torch.nn.Linear
      in_features: 64
      out_features: 32
      _input_: [metadata]

Built-in Utilities¶

Normalization¶

LazyOnlineNorm¶

Online normalization with lazy initialization and optional momentum adaptation. Normalizes inputs to zero mean and unit variance using exponential moving averages.

embedding:
  _target_: falcon.embeddings.LazyOnlineNorm
  _input_: [x]
  momentum: 0.01

DiagonalWhitener¶

Diagonal whitening with optional Hartley transform preprocessing and momentum-based running statistics. dim is required and must equal the number of features in the flattened input.

embedding:
  _target_: falcon.embeddings.DiagonalWhitener
  _input_: [x]
  dim: 64    # must match flattened input size

ToeplitzWhitener¶

Whitening for 1D signals with Toeplitz-structured covariance (e.g. stationary time series). Uses a Hartley-space EMA to track and whiten correlations efficiently without storing the full covariance matrix.

embedding:
  _target_: falcon.embeddings.ToeplitzWhitener
  _input_: [x]
  momentum: 0.1

hartley_transform¶

A standalone function for computing the discrete Hartley transform, useful as a preprocessing step.

Dimensionality Reduction¶

DynamicSVD¶

Streaming SVD with Procrustes-stabilized output and optional whitening. Maintains an eigenbasis via momentum-blended updates; Procrustes alignment keeps output coefficients stable across updates, making it safe to feed directly into a neural network.

embedding:
  _target_: falcon.embeddings.DynamicSVD
  _input_: [x]
  n_components: 32     # output dimension
  buffer_size: 128     # samples accumulated before each SVD update (default: 4 × n_components)
  momentum: 0.1        # blending weight for each SVD update

Class Reference¶

instantiate_embedding ¶

instantiate_embedding(embedding_config)

Instantiate embedding pipeline from config.

When embedding_config is None, returns a pass-through embedding that concatenates all condition tensors along the last dimension.

Source code in src/falcon/embeddings/builder.py

def instantiate_embedding(embedding_config: Dict[str, Any]) -> EmbeddingWrapper:
    """Instantiate embedding pipeline from config.

    When *embedding_config* is ``None``, returns a pass-through embedding that
    concatenates all condition tensors along the last dimension.
    """
    if embedding_config is None:
        return _PassthroughEmbedding()
    required_input_keys = _collect_input_keys(embedding_config)
    modules, input_keys_list, output_keys, _ = _flatten_config_to_modules(
        embedding_config
    )
    return EmbeddingWrapper(modules, input_keys_list, output_keys, required_input_keys)

EmbeddingWrapper ¶

EmbeddingWrapper(modules, input_keys_list, output_keys, required_input_keys)

Bases: Module

Sequential execution of modules with shared data dictionary.

Source code in src/falcon/embeddings/builder.py

def __init__(
    self,
    modules: List[nn.Module],
    input_keys_list: List[List[str]],
    output_keys: List[str],
    required_input_keys: List[str],
):
    super().__init__()
    self.modules_list = nn.ModuleList(modules)
    self.input_keys_list = input_keys_list
    self.output_keys = output_keys
    self.input_keys = required_input_keys

LazyOnlineNorm `module-attribute` ¶

LazyOnlineNorm = RunningNorm

DiagonalWhitener ¶

DiagonalWhitener(dim, momentum=0.1, eps=1e-08, use_fourier=False, track_mean=True)

Bases: Module

dim: number of features (last dimension of x) momentum: how much of the new batch stats to use (PyTorch-style) eps: small constant for numerical stability use_fourier: if True, apply Hartley transform before whitening

Source code in src/falcon/embeddings/norms.py

def __init__(self, dim, momentum=0.1, eps=1e-8, use_fourier=False, track_mean=True):
    """
    dim: number of features (last dimension of x)
    momentum: how much of the new batch stats to use (PyTorch-style)
    eps: small constant for numerical stability
    use_fourier: if True, apply Hartley transform before whitening
    """
    super().__init__()
    self.dim = dim
    self.momentum = momentum
    self.eps = eps
    self.use_fourier = use_fourier
    self.track_mean = track_mean

    self.register_buffer("running_mean", torch.zeros(dim))
    self.register_buffer("running_var", torch.ones(dim))
    self.initialized = False

update ¶

update(x)

Update running mean and variance from current batch x: Tensor of shape (batch_size, dim)

Source code in src/falcon/embeddings/norms.py

def update(self, x):
    """
    Update running mean and variance from current batch
    x: Tensor of shape (batch_size, dim)
    """
    if self.use_fourier:
        x = hartley_transform(x)

    batch_mean = x.mean(dim=0)
    batch_var = x.var(dim=0, unbiased=False)

    if not self.initialized:
        self.running_mean = batch_mean.detach()
        self.running_var = batch_var.detach()
        self.initialized = True
    else:
        self.running_mean = (
            1 - self.momentum
        ) * self.running_mean + self.momentum * batch_mean.detach()
        self.running_var = (
            1 - self.momentum
        ) * self.running_var + self.momentum * batch_var.detach()
    if not self.track_mean:
        self.running_mean *= 0

call ¶

__call__(x)

Apply whitening: (x - mean) / std If use_fourier, whitening happens in Hartley space and is transformed back.

Source code in src/falcon/embeddings/norms.py

def __call__(self, x):
    """
    Apply whitening: (x - mean) / std
    If use_fourier, whitening happens in Hartley space and is transformed back.
    """
    if self.use_fourier:
        x = hartley_transform(x)

    std = torch.sqrt(self.running_var + self.eps)
    x_white = (x - self.running_mean) / std

    if self.use_fourier:
        x_white = hartley_transform(x_white)

    return x_white

ToeplitzWhitener ¶

ToeplitzWhitener(momentum=0.1, eps=1e-08)

Bases: Module

Whitener for 1D time series assuming stationary (Toeplitz) noise covariance.

Estimates per-frequency variance via EMA in Hartley space and whitens by dividing by the estimated std. Noise is assumed zero-mean.

update(noise) — update EMA variance from a batch of noise samples call(x) — whiten x (Hartley → divide by std → inverse Hartley)

Source code in src/falcon/embeddings/norms.py

def __init__(self, momentum: float = 0.1, eps: float = 1e-8) -> None:
    super().__init__()
    self.momentum = momentum
    self.eps = eps
    self.register_buffer("running_var", None)
    self.initialized = False

update ¶

update(noise)

Update EMA variance from noise samples of shape (batch_size, T).

Source code in src/falcon/embeddings/norms.py

def update(self, noise: torch.Tensor) -> None:
    """Update EMA variance from noise samples of shape (batch_size, T)."""
    h = hartley_transform(noise)
    batch_var = h.var(dim=0, unbiased=False).detach()
    if not self.initialized:
        self.running_var = batch_var
        self.initialized = True
    else:
        self.running_var = (1 - self.momentum) * self.running_var + self.momentum * batch_var

call ¶

__call__(x)

Whiten x of shape (batch_size, T).

Source code in src/falcon/embeddings/norms.py

def __call__(self, x: torch.Tensor) -> torch.Tensor:
    """Whiten x of shape (batch_size, T)."""
    h = hartley_transform(x)
    h_white = h / torch.sqrt(self.running_var + self.eps)
    return hartley_transform(h_white)

hartley_transform ¶

hartley_transform(x)

Hartley transform: H(x) = Re(FFT(x)) - Im(FFT(x)) It is its own inverse: H(H(x)) = x

Source code in src/falcon/embeddings/norms.py

def hartley_transform(x):
    """
    Hartley transform: H(x) = Re(FFT(x)) - Im(FFT(x))
    It is its own inverse: H(H(x)) = x
    """
    fft = torch.fft.fft(x, dim=-1, norm="ortho")
    return fft.real - fft.imag

DynamicSVD ¶

DynamicSVD(n_components=10, buffer_size=None, momentum=0.1, shrinkage=True, whitener=None)

Bases: Module

Streaming SVD with Procrustes-stabilized output and optional whitening.

Maintains eigenbasis (V, Λ) via momentum-blended, eigenvalue-scaled SVD updates. Procrustes alignment ensures output coefficients have stable meaning across updates — critical when feeding into a neural network.

Optionally wraps a whitener (e.g. DiagonalWhitener). When a whitener is provided, update(x, signal) computes noise = x - signal, updates the whitener from the noise, then whitens x before the SVD update. forward() applies whitening at inference without updating statistics.

Update (when buffer full): U = [ √(1-α) · diag(√Λ_old) · V_old ; √(α/M) · X_white ] SVD(U) → V_new, Λ_new = S² Procrustes(V_new, R_old @ V_old) → R_new

forward(x) → stable k-dim coefficients: 1. c = X_white @ V.T (project onto eigenbasis) 2. c *= λ/(λ+1) (Wiener filter, diagonal) 3. c /= √λ (normalize to ~unit variance) 4. c_out = c @ R.T (rotate to stable frame)

Source code in src/falcon/embeddings/svd.py

def __init__(
    self,
    n_components: int = 10,
    buffer_size: Optional[int] = None,
    momentum: float = 0.1,
    shrinkage: bool = True,
    whitener=None,
) -> None:
    super().__init__()
    self.n_components = n_components
    self.buffer_size = buffer_size if buffer_size is not None else 4 * n_components
    self.momentum = momentum
    self.shrinkage = shrinkage
    self.whitener = whitener

    self.buffer: List[torch.Tensor] = []
    self.buffer_counter: int = 0

    self.components: Optional[torch.Tensor] = None   # (k, D)
    self.eigenvalues: Optional[torch.Tensor] = None  # (k,)
    self._R: Optional[torch.Tensor] = None           # (k, k)

update ¶

update(x, signal=None)

Accumulate a batch; trigger SVD update when buffer is full.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input data, shape (batch_size, D).	required
`signal`	`Optional[Tensor]`	True signal estimate, same shape as x. If provided and a whitener is attached, noise = x - signal is used to update the whitener before whitening x.	`None`

Source code in src/falcon/embeddings/svd.py

def update(self, x: torch.Tensor, signal: Optional[torch.Tensor] = None) -> None:
    """Accumulate a batch; trigger SVD update when buffer is full.

    Args:
        x: Input data, shape (batch_size, D).
        signal: True signal estimate, same shape as x. If provided and a
                whitener is attached, noise = x - signal is used to update
                the whitener before whitening x.
    """
    if self.whitener is not None and signal is not None:
        self.whitener.update((x - signal).detach())

    x_white = self.whitener(x) if self.whitener is not None else x

    self.buffer.append(x_white)
    self.buffer_counter += x_white.shape[0]

    if self.buffer_counter >= self.buffer_size:
        self._svd_update()
        self.buffer = []
        self.buffer_counter = 0

forward ¶

forward(x, signal=None)

Project to stable k-dimensional coefficients.

When in training mode, automatically accumulates x into the buffer and triggers an SVD update when the buffer is full. This makes DynamicSVD usable as a drop-in nn.Module embedding without a separate update() call.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input data, shape (batch_size, D).	required
`signal`	`Optional[Tensor]`	If provided and a whitener is attached, used to estimate noise for whitener updates (passed through to update()).	`None`

Returns:

Type	Description
`Tensor`	Coefficients of shape (batch_size, k), ~unit variance. Returns
`Tensor`	random noise before the first SVD update (avoids all-zero gradients).

Source code in src/falcon/embeddings/svd.py

def forward(self, x: torch.Tensor, signal: Optional[torch.Tensor] = None) -> torch.Tensor:
    """Project to stable k-dimensional coefficients.

    When in training mode, automatically accumulates x into the buffer and
    triggers an SVD update when the buffer is full. This makes DynamicSVD
    usable as a drop-in nn.Module embedding without a separate update() call.

    Args:
        x: Input data, shape (batch_size, D).
        signal: If provided and a whitener is attached, used to estimate
                noise for whitener updates (passed through to update()).

    Returns:
        Coefficients of shape (batch_size, k), ~unit variance. Returns
        random noise before the first SVD update (avoids all-zero gradients).
    """
    if self.training:
        with torch.no_grad():
            self.update(x.detach(), signal)

    if self.components is None:
        return torch.randn(x.shape[0], self.n_components, dtype=x.dtype, device=x.device)

    x_white = self.whitener(x) if self.whitener is not None else x
    c = x_white @ self.components.T

    if self.shrinkage and self.eigenvalues is not None:
        Λ = self.eigenvalues.clamp(min=1e-12)
        c = c * (Λ / (Λ + 1.0) / torch.sqrt(Λ)).unsqueeze(0)

    return c @ self._R.T

reconstruct ¶

reconstruct(x)

Wiener-filter and reconstruct in whitened D-dimensional space.

Source code in src/falcon/embeddings/svd.py

def reconstruct(self, x: torch.Tensor) -> torch.Tensor:
    """Wiener-filter and reconstruct in whitened D-dimensional space."""
    if self.components is None:
        raise ValueError("Call update() enough times before reconstruct().")

    x_white = self.whitener(x) if self.whitener is not None else x
    x_proj = x_white @ self.components.T

    if self.shrinkage:
        shrink = self.eigenvalues / (self.eigenvalues + 1.0)
        x_proj = x_proj * shrink.unsqueeze(0)

    return x_proj @ self.components

Embeddings¶

Overview¶

Declarative Configuration¶

Basic Embedding¶

Multi-Input Embedding¶

Nested Pipeline¶

Built-in Utilities¶

Normalization¶

LazyOnlineNorm¶

DiagonalWhitener¶

ToeplitzWhitener¶

hartley_transform¶

Dimensionality Reduction¶

DynamicSVD¶

Class Reference¶

instantiate_embedding ¶

EmbeddingWrapper ¶

LazyOnlineNorm module-attribute ¶

DiagonalWhitener ¶

update ¶

__call__ ¶

ToeplitzWhitener ¶

update ¶

__call__ ¶

hartley_transform ¶

DynamicSVD ¶

update ¶

forward ¶

reconstruct ¶

LazyOnlineNorm `module-attribute` ¶

call ¶

call ¶