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, wiener=True, whitener=None, fit_on_signal=False)

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 and updates the whitener from it, so the whitener normalizes the noise to unit variance. forward() applies whitening at inference without updating statistics. Without a whitener the same noise estimate is reduced to a single scalar variance (see σ² below), which is the right trade when the noise is white but not worth a per-feature model.

Inputs with more than two dimensions are flattened to (batch_size, D) internally, so image-shaped data can be fed in directly.

The signal stream (a training-only scaffold; see falcon graph configs) has two independent uses: - it is always the noise estimate for the whitener, as above; - with fit_on_signal=True it is also the stream the eigenbasis is fitted from, while projections in forward() stay on x.

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 *= λ/(λ+σ²) (Wiener filter, diagonal) 3. c /= √λ (express in units of the coefficient's own standard deviation) 4. c_out = c @ R.T (rotate to stable frame)

Steps 2-3 are gated together by wiener=True and applied as the single factor √λ/(λ+σ²). The resulting coefficient has standard deviation √(λ/(λ+σ²)) = √(SNR/(SNR+1)): order 1 for directions measured well above the noise, shrinking to 0 for those buried in it. So the output is SNR-weighted, not flatly unit-variance -- unit scale is only the λ >> σ² asymptote, and suppressing the rest is the point.

wiener=False returns the raw projection instead, whose scale is √(λ+σ²) and therefore varies by orders of magnitude across components. That is plain PCA scores; it is rarely what you want to hand to a network.

reconstruct() gates on the same flag but applies only step 2, since it returns to data space where step 3's rescaling would be wrong.

σ² is the noise variance in the units the eigenvalues are measured in: - with a whitener, the noise is normalized per feature, so σ² = 1; - without one, a single scalar is estimated from x - signal in update() under a white-noise assumption, so the denominator stays commensurate with λ instead of comparing raw signal power against an arbitrary 1; - with neither a whitener nor a signal, nothing can be estimated and it falls back to 1.

Step 2 is the textbook Wiener gain only when λ is signal power, i.e. when fit_on_signal=True. Fitting on x instead gives λ ≈ λ_signal + σ², so the gain becomes (λ_s+σ²)/(λ_s+2σ²) >= 1/2 and a noise-only direction is passed through at half amplitude rather than suppressed.

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

def __init__(
    self,
    n_components: int = 10,
    buffer_size: Optional[int] = None,
    momentum: float = 0.1,
    wiener: bool = True,
    whitener=None,
    fit_on_signal: bool = False,
) -> None:
    super().__init__()
    self.fit_on_signal = fit_on_signal
    self.n_components = n_components
    self.buffer_size = buffer_size if buffer_size is not None else 4 * n_components
    self.momentum = momentum
    self.wiener = wiener
    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)

    # Scalar noise variance used in the Wiener denominator.  None means
    # "assume 1", which is correct when a whitener is attached (it
    # normalizes the noise) and is the fallback when nothing better is
    # known.  Estimated from `signal` in update() otherwise.
    self._noise_var: Optional[torch.Tensor] = None

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) or (batch_size, ...).	required
signal	Optional[Tensor]	True signal estimate, same shape as x. When provided, noise = x - signal updates the whitener if one is attached, or the scalar noise-variance estimate if not. With fit_on_signal=True it is additionally what the eigenbasis is fitted from, in place of 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) or (batch_size, ...).
        signal: True signal estimate, same shape as x. When provided,
                noise = x - signal updates the whitener if one is attached,
                or the scalar noise-variance estimate if not. With
                fit_on_signal=True it is additionally what the eigenbasis is
                fitted from, in place of x.
    """
    if x.dim() > 2:
        x = x.flatten(start_dim=1)
    if signal is not None and signal.dim() > 2:
        signal = signal.flatten(start_dim=1)
    if signal is not None:
        noise = (x - signal).detach()
        if self.whitener is not None:
            # Whitener normalizes the noise per feature, so downstream the
            # noise variance is 1 by construction and _noise_var stays None.
            self.whitener.update(noise)
        else:
            # No whitener: estimate a single scalar noise variance under a
            # white-noise assumption, so the Wiener denominator below is in
            # the same units as the eigenvalues instead of an arbitrary 1.
            batch_var = noise.var(unbiased=False)
            self._noise_var = batch_var if self._noise_var is None else (
                (1.0 - self.momentum) * self._noise_var
                + self.momentum * batch_var
            )

    # fit_on_signal: the eigenbasis is fitted from the noise-free ``signal``
    # stream (a training-only scaffold); projections in forward() remain on
    # ``x``. Removes noise-contaminated components (empirical eigenvalues up
    # to the Marchenko-Pastur edge (1+sqrt(D/N))^2 sigma^2 masquerade as
    # structure when fitting on noisy x). With wiener=True the trailing
    # near-zero clean eigenvalues are damped, so the 1/sqrt(lambda)
    # rescaling cannot amplify observation noise.
    fit_x = signal if (self.fit_on_signal and signal is not None) else x
    x_white = self.whitener(fit_x) if self.whitener is not None else fit_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]	Passed through to update() during training, where it feeds the whitener's noise estimate and, with fit_on_signal=True, the eigenbasis fit. Never used for the projection itself, so it is not needed at inference.	None

Returns:

Type	Description
Tensor	Coefficients of shape (batch_size, k). With wiener=True each has
Tensor	standard deviation √(SNR/(SNR+1)) -- order 1 where the direction is
Tensor	well measured, near 0 where it is noise-dominated. Returns random
Tensor	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: Passed through to update() during training, where it feeds
                the whitener's noise estimate and, with fit_on_signal=True,
                the eigenbasis fit.  Never used for the projection itself,
                so it is not needed at inference.

    Returns:
        Coefficients of shape (batch_size, k).  With wiener=True each has
        standard deviation √(SNR/(SNR+1)) -- order 1 where the direction is
        well measured, near 0 where it is noise-dominated.  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 x.dim() > 2:
        x = x.flatten(start_dim=1)

    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.wiener and self.eigenvalues is not None:
        Λ = self.eigenvalues.clamp(min=1e-12)
        # sqrt(Λ)/(Λ+σ²) == Λ/(Λ+σ²) · 1/sqrt(Λ), i.e. the Wiener gain and
        # the unit-variance normalization folded into one factor.  Written
        # this way it avoids dividing by sqrt of the 1e-12 clamp.
        c = c * (torch.sqrt(Λ) / (Λ + self._sigma2())).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().")

    if x.dim() > 2:
        x = x.flatten(start_dim=1)

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

    if self.wiener:
        shrink = self.eigenvalues / (self.eigenvalues + self._sigma2())
        x_proj = x_proj * shrink.unsqueeze(0)

    return x_proj @ self.components