mixle.models.quotient module

A translation-quotient leaf: conv feature map -> global pool -> softmax, as a mixle conditional-density leaf.

TranslationQuotientLeaf(module) wraps a Torch module whose forward pass factors as conv_stack -> global_pool -> linear and declares its symmetry group via leaf.group == "translation". Global average/max pooling after a conv stack (with same-padding convolutions) makes the module’s output exactly invariant to integer pixel shifts up to the boundary effect of zero-padding: shifting the input by (dy, dx) shifts the conv feature maps by the same amount, and the pool discards spatial position entirely – so log_density(x) == log_density(shift(x)) up to whatever boundary pixels the shift dragged in/out of the receptive field. This module also builds UnpooledConvLeaf, a same-capacity baseline (matching conv depth/width, flatten + dense, no pooling) for an apples-to-apples comparison of the “quotient” (pooled, group-invariant) leaf against an unstructured head with the same feature extractor.

Follows the declare-a-leaf/fit-via-optimize() pattern used elsewhere in mixle.models (see mixle.models.softmax_leaf.NeuralCategorical) rather than a bespoke torch loop: both leaves here are thin torch.nn.Module builders plus a group/declared_group() tag; fitting goes through NeuralCategorical(module).estimator() and mixle.inference.optimize exactly like any other softmax leaf.

Requires torch. Treat this as an experimental modeling option and compare it against the unpooled baseline before making a release claim about benefit.

conv_feature_stack(in_channels=3, hidden_channels=16, out_channels=32)[source]

A small two-layer same-padding conv feature extractor, shared by both leaves below.

Same-padding (padding=1 for 3x3 kernels) keeps spatial shifts of the input exactly reflected as spatial shifts of the output feature map (away from the boundary), which is what makes the pooled leaf’s translation invariance hold by construction.

Parameters:
  • in_channels (int)

  • hidden_channels (int)

  • out_channels (int)

Return type:

Any

build_translation_quotient_module(n_classes, in_channels=3, hidden_channels=16, out_channels=32)[source]

Build the quotient module: conv stack -> global average pool -> linear -> logits.

Global pooling erases spatial position from the feature map entirely, so the classifier head sees the same input (up to boundary truncation) regardless of where the pattern sits in the image – the “quotient by the translation group” this leaf is named for.

Parameters:
  • n_classes (int)

  • in_channels (int)

  • hidden_channels (int)

  • out_channels (int)

Return type:

Any

build_unpooled_conv_module(n_classes, spatial_size, in_channels=3, hidden_channels=16, out_channels=32)[source]

Build the same-capacity baseline module: the identical conv stack, but flatten + dense (no pooling).

Same conv depth/width as build_translation_quotient_module() so the comparison isolates the effect of the pooling/quotient step rather than differences in feature-extractor capacity.

Parameters:
  • n_classes (int)

  • spatial_size (int)

  • in_channels (int)

  • hidden_channels (int)

  • out_channels (int)

Return type:

Any

class TranslationQuotientLeaf(module, **neural_categorical_kwargs)[source]

Bases: object

p(y | x) = softmax(module(x)) for a conv->global-pool module, declaring the “translation” group.

Thin wrapper around mixle.models.softmax_leaf.NeuralCategorical that adds the group-declaration part of the leaf contract: leaf.group == "translation" (also exposed as leaf.declared_group() for callers that prefer a method). Fitting/serialization/log-density all delegate to the wrapped NeuralCategorical – this class does not reimplement the leaf contract, it just tags a NeuralCategorical built from a pooled conv module with its symmetry group.

Parameters:
  • module (Any)

  • neural_categorical_kwargs (Any)

declared_group()[source]

Return the symmetry group this leaf’s density is invariant to.

Return type:

str

log_density(xy)[source]

Delegate log p(y | x) scoring to the wrapped neural-categorical leaf.

Parameters:

xy (Any)

Return type:

float

seq_log_density(enc)[source]

Delegate vectorized conditional log-probability scoring to the wrapped leaf.

Parameters:

enc (Any)

Return type:

Any

predict(x)[source]

Return class predictions from the wrapped neural-categorical leaf.

Parameters:

x (Any)

Return type:

Any

estimator(pseudo_count=None)[source]

Return the wrapped leaf’s estimator.

Parameters:

pseudo_count (float | None)

Return type:

Any

sampler(seed=None)[source]

Return the wrapped leaf’s conditional sampler.

Parameters:

seed (int | None)

Return type:

Any

class UnpooledConvLeaf(module, **neural_categorical_kwargs)[source]

Bases: object

Same-capacity baseline: p(y | x) = softmax(module(x)) for a conv->flatten->dense module.

No symmetry group is declared (group is None) – this baseline has no built-in translation invariance, which is exactly the property TranslationQuotientLeaf is compared against.

Parameters:
  • module (Any)

  • neural_categorical_kwargs (Any)

declared_group()[source]

Return None because this baseline declares no invariance group.

Return type:

str | None

log_density(xy)[source]

Delegate log p(y | x) scoring to the wrapped neural-categorical leaf.

Parameters:

xy (Any)

Return type:

float

seq_log_density(enc)[source]

Delegate vectorized conditional log-probability scoring to the wrapped leaf.

Parameters:

enc (Any)

Return type:

Any

predict(x)[source]

Return class predictions from the wrapped neural-categorical leaf.

Parameters:

x (Any)

Return type:

Any

estimator(pseudo_count=None)[source]

Return the wrapped leaf’s estimator.

Parameters:

pseudo_count (float | None)

Return type:

Any

sampler(seed=None)[source]

Return the wrapped leaf’s conditional sampler.

Parameters:

seed (int | None)

Return type:

Any

shift_image_batch(x, dy, dx)[source]

Zero-pad shift an (n, c, h, w) batch by (dy, dx) pixels (numpy in, numpy out).

Used to build the “corrupted” (shifted) test set for the robustness comparison and to test the invariance property: pixels shifted out of frame are dropped, pixels shifted into frame are zero.

Parameters:
Return type:

Any