Weighted sums for a requested group

Select an interleaved category, bind its item values and weights, and predict that group’s weighted sum.

Seed 3605 on gpu: 5 of 5 behavioral checks met. See the measurements below. The experiment code has changed since this run; these measurements describe its earlier version.

Insights

The model can select a requested group and learn its weighted total from raw values and weights. The same bag appears with different requested groups, so bag content alone cannot determine the answer. Rotating group labels tests selection; swapping weights within each group separately tests the value–weight relationship. Both controls remove accuracy against retained original labels.

Keeping group, value, and weight together preserves the relationships the answer needs. The supplied-contribution control helps separate multiplication from selection and reduction. The result is still narrow: groups are familiar and equally populated. Unseen labels, absent groups, and varying group sizes are additional tasks, not consequences of this pass.

Setup

"""The natural schema should learn a selected group's weighted sum.

Learn filtering, coordinate binding, multiplication, and summation.

Run: uv run python proofs/run.py P010"""

from __future__ import annotations

from collections.abc import Iterator
from copy import deepcopy

import lightning.pytorch as lit
import numpy as np
import torch
from reporting import report

import relflow as rf

PROOF_ID = "P010"
GROUPS = ("A", "B", "C")
ITEMS_PER_GROUP = 2
ITEMS = len(GROUPS) * ITEMS_PER_GROUP

Examples

Targets use mask=True; the labels below are supervision hidden from the encoder.

Request group A

selected_group: A
items:
  - {group: B, value: 0.3, weight: 0.4}
  - {group: A, value: 0.8, weight: 0.5}
  - {group: C, value: -0.5, weight: 1.4}
  - {group: A, value: -0.4, weight: 1.25}
  - {group: B, value: -0.2, weight: 0.9}
  - {group: C, value: 0.7, weight: 0.3}
answer: -0.1

Only the two A items contribute: 0.4 − 0.5 = −0.1.

Request another group from the same bag

selected_group: B
items:
  - {group: B, value: 0.3, weight: 0.4}
  - {group: A, value: 0.8, weight: 0.5}
  - {group: C, value: -0.5, weight: 1.4}
  - {group: A, value: -0.4, weight: 1.25}
  - {group: B, value: -0.2, weight: 0.9}
  - {group: C, value: 0.7, weight: 0.3}
answer: -0.06

The unchanged bag now answers for B: 0.12 − 0.18 = −0.06. The request is necessary because each training bag appears once for every group.

Swap weights within each group

selected_group: A
items:
  - {group: B, value: 0.3, weight: 0.9}
  - {group: A, value: 0.8, weight: 1.25}
  - {group: C, value: -0.5, weight: 0.3}
  - {group: A, value: -0.4, weight: 0.5}
  - {group: B, value: -0.2, weight: 0.4}
  - {group: C, value: 0.7, weight: 1.4}
answer: -0.1  # Retained original label

The A pairs now imply 0.8. This control deliberately retains −0.1 as the original target; every group keeps its values, weights, and membership counts, but the local pairings change.

Synthetic data and controls

def records(*, bags: int, seed: int) -> Iterator[dict]:
    """Draw interleaved groups and emit one request for every selected group."""
    rng = np.random.default_rng(seed)
    for bag in range(bags):
        items: list[dict[str, object]] = []
        answers: dict[str, float] = {}
        for group in GROUPS:
            values = rng.uniform(-1.0, 1.0, size=ITEMS_PER_GROUP)
            weights = rng.uniform(0.2, 1.8, size=ITEMS_PER_GROUP)
            contributions = values * weights
            answers[group] = float(contributions.sum())
            for value, weight, contribution in zip(values, weights, contributions, strict=True):
                item: dict[str, object] = {"group": group}
                item.update(value=float(value), weight=float(weight))
                items.append(item)
        interleaved = [items[index] for index in rng.permutation(len(items))]
        for selected_group in GROUPS:
            yield {
                "bag": bag,
                "selected_group": selected_group,
                "items": interleaved,
                "answer": answers[selected_group],
            }


def rotate_group_labels(observations: list[dict]) -> list[dict]:
    """Break group/item association while preserving all group marginals."""
    successor = {group: GROUPS[(index + 1) % len(GROUPS)] for index, group in enumerate(GROUPS)}
    rows = deepcopy(observations)
    for row in rows:
        row["items"] = [{**item, "group": successor[item["group"]]} for item in row["items"]]
    return rows


def swap_weights_within_groups(observations: list[dict]) -> list[dict]:
    """Break value/weight pairing without changing any per-group marginal."""
    rows = deepcopy(observations)
    cached: dict[int, list[dict[str, object]]] = {}
    for row in rows:
        bag = int(row["bag"])
        if bag not in cached:
            items = [dict(item) for item in row["items"]]
            for group in GROUPS:
                indices = [index for index, item in enumerate(items) if item["group"] == group]
                if len(indices) != ITEMS_PER_GROUP:
                    raise ValueError(f"bag {bag} has {len(indices)} items in group {group!r}")
                weights = [items[index]["weight"] for index in indices]
                for index, weight in zip(indices, np.roll(weights, 1), strict=True):
                    items[index]["weight"] = float(weight)
            cached[bag] = items
        row["items"] = cached[bag]
    return rows


def permute_items(observations: list[dict], *, seed: int) -> list[dict]:
    """Jointly permute complete items while retaining targets."""
    rng = np.random.default_rng(seed)
    rows = deepcopy(observations)
    cached: dict[int, list[dict[str, object]]] = {}
    for row in rows:
        bag = int(row["bag"])
        if bag not in cached:
            items = row["items"]
            cached[bag] = [items[index] for index in rng.permutation(len(items))]
        row["items"] = cached[bag]
    return rows


def prediction(model: rf.Model, observations: list[dict]) -> np.ndarray:
    inputs = [{"selected_group": row["selected_group"], "items": row["items"]} for row in observations]
    output = model.predict(inputs)["predictions"].to_pylist()
    return np.asarray([row["request/answer"]["content"] for row in output], dtype=np.float64)


def rmse(actual: np.ndarray, predicted: np.ndarray | float) -> float:
    return float(np.sqrt(np.mean(np.square(actual - predicted))))


def score(*, train: list[dict], test: list[dict], predicted: np.ndarray) -> dict[str, float]:
    """Compare held-out RMSE with the constant training-target mean."""
    actual = np.asarray([row["answer"] for row in test], dtype=np.float64)
    baseline_rmse = rmse(actual, float(np.asarray([row["answer"] for row in train], dtype=np.float64).mean()))
    measured = rmse(actual, predicted)
    return {"rmse": measured, "baseline_rmse": baseline_rmse, "nrmse": measured / baseline_rmse}

Model tree

Figure 1: The item branch keeps all tokens; root Attention compresses them into six outputs for the selected-group answer.

How it works

Each bag contains two items from each of A, B, and C. The generator emits the same bag once per requested group, so the bag alone cannot determine the answer. reduction=None preserves the item branch’s encoded coordinates; root rf.Attention combines them with the visible selected_group request.

Rotating item group labels tests group selection. Swapping weights within each group independently tests value–weight binding while preserving per-group marginals. Both corruptions keep original labels and must worsen error. Reordering complete items should preserve predictions. The supplied-product control separates multiplication from selection and reduction.

Training and evaluation

def run(seed: int, steps: int | None, accelerator: str) -> tuple[dict, dict]:
    lit.seed_everything(seed, workers=True)
    # Split seeds are independent; rerunning a generator reproduces the same records.
    train = list(records(bags=512, seed=seed + 1))
    test = list(records(bags=256, seed=seed + 3))
    model = rf.Model(
        name="request",
        d_model=48,
        n_layers=2,
        n_heads=4,
        reduction=rf.Attention(n_outputs=ITEMS, n_layers=2),
        batch_size=96,
        optimizer=lambda module: torch.optim.Adam(module.parameters(), lr=0.001),
        items=rf.Branch(
            length=ITEMS,
            n_layers=2,
            reduction=None,
            group=rf.Category(size=len(GROUPS), p_unavailable=0.0),
            value=rf.Number,
            weight=rf.Number,
        ),
        selected_group=rf.Category(size=len(GROUPS), p_unavailable=0.0),
        answer=rf.Number(mask=True, objective="mse"),
    )
    datamodule = rf.SyntheticDataModule(
        model=model,
        train=lambda: records(bags=512, seed=seed + 1),
        validate=lambda: records(bags=128, seed=seed + 2),
        seed=seed,
    )
    trainer = lit.Trainer(
        accelerator=accelerator,
        devices=1,
        max_epochs=-1,
        max_steps=1100 if steps is None else min(steps, 1100),
        logger=False,
        enable_progress_bar=False,
        enable_model_summary=False,
        enable_checkpointing=False,
        deterministic=True,
        num_sanity_val_steps=0,
    )
    trainer.fit(model, datamodule=datamodule)

    # Evaluate held-out answers and retain their original labels in corruption controls.
    intact_prediction = prediction(model, test)
    intact = score(train=train, test=test, predicted=intact_prediction)
    labels = score(train=train, test=test, predicted=prediction(model, rotate_group_labels(test)))
    pairing = score(train=train, test=test, predicted=prediction(model, swap_weights_within_groups(test)))
    permuted_prediction = prediction(model, permute_items(test, seed=seed + 4))
    target_scale = float(np.std(np.asarray([row["answer"] for row in test], dtype=np.float64)))
    permutation_drift = rmse(intact_prediction, permuted_prediction) / target_scale
    metrics = {
        "intact": intact,
        "labels": labels,
        "pairing": pairing,
        "permutation_drift": permutation_drift,
        "steps": trainer.global_step,
    }
    checks = {
        "All measurements finite": bool(
            all(
                (
                    np.isfinite(value)
                    for value in (intact["nrmse"], labels["nrmse"], pairing["nrmse"], permutation_drift)
                )
            )
        ),
        "Raw grouped sum nRMSE below 0.50": bool(intact["nrmse"] < 0.5),
        "Rotated labels raise nRMSE by at least 0.20": bool(labels["nrmse"] >= intact["nrmse"] + 0.2),
        "Swapped weights raise nRMSE by at least 0.20": bool(pairing["nrmse"] >= intact["nrmse"] + 0.2),
        "Permutation drift below 0.10 target SD": bool(permutation_drift < 0.1),
    }
    return metrics, checks

Evidence

Latest full run

Seed 3605, gpu, recorded 2026-09-15T02:20:08.067014+00:00. Outcome: met.

Source fingerprint: 27937931b9c7d89f9cced00722a005874718f90849eb68704f32bb43c221fa0d. Python 3.12.6; Torch 2.12.0.

Measurement	Value
intact	rmse: 0.0432588; baseline_rmse: 0.885622; nrmse: 0.0488456
labels	rmse: 1.26192; baseline_rmse: 0.885622; nrmse: 1.42489
pairing	rmse: 0.535116; baseline_rmse: 0.885622; nrmse: 0.604226
permutation_drift	0.0237862
steps	1100

Behavioral checks

Behavioral check	Outcome
All measurements finite	Met
Raw grouped sum nRMSE below 0.50	Met
Rotated labels raise nRMSE by at least 0.20	Met
Swapped weights raise nRMSE by at least 0.20	Met
Permutation drift below 0.10 target SD	Met

Recorded results.

Remaining work

Repeat all gates across seeds, then vary group sizes and test an absent requested group. Unseen labels, simultaneous requests, weight scaling, and selected-group weighted means remain open. These scores are not a capacity-matched comparison with the supplied-product control.

The family’s promotion target is at least three core seeds and ten lightweight calibration seeds.

Reproduce

Run by stable ID from the repository root:

uv run python proofs/run.py P010

Or run the self-contained script directly:

PYTHONPATH=proofs uv run python proofs/aggregation/selected_group_raw_weighted_sum.py

Add --accelerator gpu for CUDA or --seed 42 for another seeded experiment. --steps 2 checks execution with a short training budget; it is recorded as a smoke run.

Download the complete proof.

if __name__ == "__main__":
    report(PROOF_ID, run, seed=3605)