Come up with different masking methods

[kutayeroglu]

Here's output with conversation with Opus 4.7

Good question. Lines 129–130 are exactly the lever to pull if you want to keep multiblock's solid rectangles but let the NPZ noise library shape where the rectangles land, rather than (as multinoise does) shape what's inside them. Here are concrete, mutually compatible integration ideas, from minimal to deeper.

---

Background: what's available from the NPZ library

From multinoise.py, the relevant ingredients are reusable verbatim:

        # -- Color Noise Initialization
        self.color_mask_ratio = color_mask_ratio
        for name, val in [("enc_drop_order", enc_drop_order), ("pred_drop_order", pred_drop_order)]:
            if val not in ("lowest", "highest"):
                raise ValueError(f"{name} must be 'lowest' or 'highest', got {val!r}")
        self.enc_drop_order = enc_drop_order
        self.pred_drop_order = pred_drop_order
        self.trans_sequence = transforms.Compose([
            transforms.RandomCrop(self.height), # Crop to [self.height, self.width] which is [14, 14] for ViT/14
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.RandomVerticalFlip(p=0.5),
            NormalizeBySliceMax()
        ])
        self._load_color_pattern(color_noise_path)

    def _extract_noise_windows(self, B: int) -> torch.Tensor:
        """
        Produce B augmented noise grids from the stored noise pattern collection.
        ...
        Returns:
            Tensor of shape [B, height, width].
        """

So a per-batch [B, H, W] noise field is essentially free to obtain inside multiblock.MaskCollator.

---

Option A — Noise-weighted top-left sampling (smallest change to lines 129–130)

Instead of uniform sampling of the top-left corner, draw it from a categorical distribution whose logits depend on the noise grid. Two natural variants:

A.1 Score = noise at the top-left pixel

# replace the two uniform draws around lines 129–130
valid = noise_grid[: self.height - h, : self.width - w]   # [H', W']
probs = valid.flatten().clamp_min(1e-8)
probs = probs / probs.sum()
flat_idx = torch.multinomial(probs, 1).item()
top  = flat_idx //  valid.shape[1]
left = flat_idx %   valid.shape[1]

A.2 Score = sum of noise inside the would-be box (true "energy" of placement)

Pre-compute an integral image of the noise grid once per (B, image) and read box sums in O(1):

ii = noise_grid.cumsum(0).cumsum(1)   # integral image; do once per image
def box_sum(t, l):
    return ii[t+h-1, l+w-1] - (ii[t-1, l+w-1] if t else 0) \
                            - (ii[t+h-1, l-1] if l else 0) \
                            + (ii[t-1, l-1] if t and l else 0)
scores = torch.stack([
    box_sum(t, l)
    for t in range(self.height - h)
    for l in range(self.width  - w)
])

Then torch.multinomial(softmax(scores / T), 1) (with a temperature T) gives you a smooth control between "uniform" (T → ∞) and "always pick the noisiest region" (T → 0). You can also flip the sign of scores to bias toward low-noise regions instead.

Effect: rectangles still solid, but they preferentially land on (or away from) noise-rich locations. This gives you a structured curriculum of where to predict without the irregularity that multinoise introduces.

---

Option B — Use noise to derive acceptable_regions

multiblock already has acceptable_regions-driven rejection sampling. You can plug a thresholded noise mask in as a permanent acceptable region, requiring blocks to fall on patches in (say) the top 50% of the noise field:

threshold = torch.quantile(noise_grid, 0.5)
noise_region = (noise_grid >= threshold).to(torch.int32)   # [H, W]
# pass this as the FIRST entry of acceptable_regions, before mask complements

This integrates cleanly with multiblock's existing relaxation loop:

                tries: Number of constraint relaxations (0 = enforce all, 1 = ignore one, etc.)
            
            Process:
                - acceptable_regions is a list of 2D binary masks (1 = acceptable, 0 = not acceptable)
                - Element-wise multiplication: mask *= region zeros out pixels where region=0
                - This "crops" the block to only overlap with acceptable regions
                - As tries increases, fewer regions are enforced (gradual constraint relaxation)
            """            
            N = max(int(len(acceptable_regions) - tries), 0)
            for k in range(N):
                mask *= acceptable_regions[k]  # Element-wise: 1*1=1 (keep), 1*0=0 (remove)

Caveat: this breaks the "solid rectangle" property the same way multinoise does — the constraint multiplication can carve holes. If you want to preserve solidity, do not add the noise mask to acceptable_regions; only use it to bias the top-left draw (Option A) and reject samples whose box sum is below some quantile (Option C).

---

Option C — Rejection sampling against a noise score

Stay almost identical to current multiblock lines 129–130, but reject draws whose box sum is below a percentile of all possible box sums:

top  = torch.randint(0, self.height - h, (1,))
left = torch.randint(0, self.width  - w, (1,))
if box_sum(top, left) < threshold:
    continue   # resample

Cheap, preserves rectangles, and uses the noise field as a soft spatial prior. With a slowly-rising threshold across training, you get a built-in curriculum.

---

Option D — Noise-conditioned block size or count

Instead of touching lines 129–130, use the per-image noise statistics to:

• Modulate mask_scale (currently sampled by _sample_block_size): images with high noise variance → larger predictor blocks.
• Modulate npred: spawn an extra predictor block on images whose noise field has a strong second mode.
• Choose between enc / pred drop orders dynamically per image.

These don't change geometry, only schedule, and so are very safe to A/B against vanilla multiblock.

---

Option E — Noise to break ties between candidate top-lefts in acceptable_regions

When acceptable_regions produces many feasible (top, left) positions, currently the loop just keeps trying random corners. You can replace that random retry with: enumerate feasible corners (boolean grid of valid placements), score them by box_sum against the noise grid, sample one. This kills the timeout/relaxation loop for the common case and gives noise-aware enc/pred separation without dropping any patches.

---

Recommendation

If your goal is "use the NPZ library, but keep mask blocks solid":

1. Wire up the existing noise infra (NormalizeBySliceMax, _load_color_pattern, _extract_noise_windows) into multiblock.MaskCollator.__init__ and __call__ — no logic changes inside the sampling loop yet. This is purely plumbing and lets you log the noise grids alongside masks.
2. Replace lines 129–130 with Option A.2 (integral-image-weighted top-left sampling), with a temperature hyper-parameter noise_temp that defaults to +∞ (uniform — bit-identical to current behavior). That makes the change a strict superset of vanilla multiblock.
3. Optionally layer Option C on top for a stronger prior, controlled by a single quantile knob.

That gives you a meaningful, novel use of the noise library that is orthogonal to multinoise.py: multinoise modulates what is masked inside a box; this would modulate where the box lives, while keeping the rectangles solid.

If you want, I can sketch the exact diff for Option A.2 + the plumbing in step 1 (still in Ask mode — sketch only, no edits).

[kutayeroglu]

What "uniform" means in torch.randint

The docs phrase "generated uniformly" is a statement about the probability distribution of every single call, not only about what happens in the limit. Specifically, for torch.randint(0, N, (1,)):

$$P(X = k) = \frac{1}{N} \quad \text{for every } k \in {0, 1, \ldots, N-1}$$

Every integer in [0, N) has equal probability on every draw. They are i.i.d. samples from a discrete uniform distribution.

Your statement — "as the number of draws → ∞, the empirical histogram looks uniform" — is a true consequence of this (the Law of Large Numbers), but it's a weaker description than reality:

• Per draw: each value has probability 1/N.
• In the limit of many draws: empirical frequencies converge to 1/N.

If only the second were true, you could imagine a sampler that cycles 0, 1, 2, ..., N-1, 0, 1, .... That also has a uniform empirical limit, but is not how torch.randint works — torch.randint is memoryless and i.i.d.

---

Tying this back to lines 129–130

            top = torch.randint(0, self.height - h, (1,))  # maximum valid top position
            left = torch.randint(0, self.width - w, (1,))  # maximum valid left position

For ViT-B/16 at 224×224, self.height = self.width = 14. Suppose the sampled block size is (h, w) = (4, 6). Then:

• top is drawn uniformly from {0, 1, ..., 9} (10 values).
• left is drawn uniformly from {0, 1, ..., 7} (8 values).
• The pair (top, left) is uniformly distributed over the 10 × 8 = 80 valid corner positions, each with probability 1/80.

In picture form, every "X" below is an equally-likely top-left corner of the block:

top-left grid (only positions where the 4×6 block still fits in the 14×14 patch grid)
left →   0 1 2 3 4 5 6 7 . . . . . .
top 0 :  X X X X X X X X . . . . . .
top 1 :  X X X X X X X X . . . . . .
...
top 9 :  X X X X X X X X . . . . . .
top 10:  . . . . . . . . . . . . . .   ← would put block past the bottom edge
...

So "uniform sampling of the top-left corner" literally means: every one of those 80 X's has the same 1/80 chance of being chosen this iteration, and that choice has nothing to do with image content, noise, or any prior draw.

---

Why Option A changes this

Option A replaces that uniform 1/80 with a non-uniform categorical distribution whose probabilities come from the noise grid. Concretely:

1. Build a score s(t, l) for every valid corner — either the noise value at (t, l), or the sum of noise values inside the box [t, t+h) × [l, l+w).
2. Normalize: p(t, l) = s(t, l) / Σ s(·, ·).
3. Sample one corner from this distribution with torch.multinomial.

Now corners overlapping noise-rich regions get probability higher than 1/80, and corners over noise-poor regions get probability lower than 1/80 (possibly near zero). The mechanism is still i.i.d. categorical sampling — just with non-equal weights — so:

• Per draw: corner (t, l) has probability p(t, l).
• In the limit of many draws across many images: the empirical histogram of chosen corners matches p(t, l), which mirrors the noise field.

The temperature knob I mentioned (p ∝ exp(s / T)) just controls how "peaked" p is:

• T → ∞: p flattens out → recovers the uniform 1/80 and reproduces current behavior exactly.
• T → 0: p collapses onto the single highest-noise corner → deterministic, always the same spot for a given noise grid.

So Option A is best read as: "keep the i.i.d. corner-sampling structure of torch.randint, but replace the implicit uniform weights with weights that depend on the noise field."

torch.randint docs

[kutayeroglu]

Leaving full plan here as documentation/backup

---

name: noise-guided multiblock collator
overview: Add a new ng_multiblock mask collator that keeps multiblock's solid rectangles but draws each block's top-left corner from a noise-weighted distribution (Option A). Wire it through src/train.py and ship a sample config.
todos:

• id: create_collator
  content: "Create src/masks/ng_multiblock.py: copy noise helpers from multinoise.py, add _score_grid (boxsum + corner), and replace lines 129-130 of multiblock.py with noise-weighted torch.multinomial sampling gated by per-role bias and temperature."
  status: completed
• id: wire_train
  content: "Wire ng_multiblock into src/train.py: add import, read score_mode/enc_bias/pred_bias/noise_temperature from mask config, and add an elif mask_type == 'ng_multiblock' branch constructing NGMBMaskCollator."
  status: completed
• id: sample_config
  content: "Add configs/bal_ngmb_vitb.yaml mirroring bal_mb_vitb.yaml with mask_type: ng_multiblock plus the new noise hyperparameters and write_tag bal_ngmb_vitb."
  status: completed
  isProject: false

---

1. New collator: src/masks/ng_multiblock.py

Skeleton mirrors src/masks/multiblock.py but adds the noise machinery from src/masks/multinoise.py. Key changes vs multiblock:

• Copy NormalizeBySliceMax, _load_color_pattern, _extract_noise_windows from multinoise.py verbatim (no behavioral change to those helpers).
• Extend __init__ with: color_noise_path, score_mode in {"boxsum","corner"} (default "boxsum"), enc_bias in {"high","low","none"} (default "high"), pred_bias (default "high"), noise_temperature: float (default 0.5). Validate enums in __init__.
• Add helper _score_grid(noise_grid, h, w, mode) -> Tensor[H-h, W-w]:
  • mode == "corner": return noise_grid[: H-h, : W-w] directly.
  • mode == "boxsum": build integral image ii = noise_grid.cumsum(0).cumsum(1) once, then vectorize the box-sum formula ii[t+h-1,l+w-1] - ii[t-1,l+w-1] - ii[t+h-1,l-1] + ii[t-1,l-1] over all valid (t,l) via tensor slicing.
• Replace _sample_block_mask with a version that takes noise_grid and bias arguments and replaces only src/masks/multiblock.py:129-130:
  • If bias == "none": keep the original two torch.randint calls (bit-identical fallback to vanilla multiblock).
  • Else: compute score = _score_grid(noise_grid, h, w, score_mode); min-max normalize to [0,1]; if bias == "low", replace score with 1 - score; compute probs = softmax(score.flatten() / T); draw a single index with torch.multinomial and decode back to (top, left).
  • Everything downstream (rectangle fill, acceptable_regions constraint loop, min_keep rejection, complement) stays exactly as in multiblock.py. mask_complement remains the rectangle complement so the existing enc/pred separation logic is unchanged.
• __call__: identical control flow to multiblock.MaskCollator.__call__, with two additions:
  • After e_size/p_size, call batch_noise_grids = self._extract_noise_windows(B) (shape [B, H, W]).
  • Per image i, pass noise_grid=batch_noise_grids[i] plus the role-specific bias (pred_bias for pred blocks, enc_bias for enc blocks) into _sample_block_mask.

2. Wiring in src/train.py

• Add import alongside the others at src/train.py:32-34:

  from src.masks.ng_multiblock import MaskCollator as NGMBMaskCollator

• Update the mask_type doc comment at src/train.py:115 to include ng_multiblock.
• After src/train.py:119 read the new args:

  score_mode = args["mask"].get("score_mode", "boxsum")
  enc_bias  = args["mask"].get("enc_bias",  "high")
  pred_bias = args["mask"].get("pred_bias", "high")
  noise_temperature = float(args["mask"].get("noise_temperature", 0.5))

• Add a new branch in the dispatch at src/train.py:329-381, before the else for vanilla multiblock:

  elif mask_type == "ng_multiblock":
      if green_noise_data_path is None:
          raise ValueError("green_noise_data_path must be specified when using ng_multiblock mask type")
      mask_collator = NGMBMaskCollator(
          input_size=crop_size, patch_size=patch_size,
          enc_mask_scale=enc_mask_scale, pred_mask_scale=pred_mask_scale,
          aspect_ratio=aspect_ratio,
          nenc=num_enc_masks, npred=num_pred_masks,
          min_keep=min_keep, allow_overlap=allow_overlap,
          debug_log=os.environ.get("LOG_MULTIBLOCK_DEBUG", "") == "1",
          color_noise_path=green_noise_data_path,
          score_mode=score_mode,
          enc_bias=enc_bias, pred_bias=pred_bias,
          noise_temperature=noise_temperature,
      )

3. Sample config: configs/bal_ngmb_vitb.yaml

Clone configs/bal_mb_vitb.yaml and override the mask: block plus write_tag:

mask:
  allow_overlap: false
  aspect_ratio: [0.75, 1.5]
  enc_mask_scale: [0.85, 1.0]
  min_keep: 10
  num_enc_masks: 1
  num_pred_masks: 4
  patch_size: 16
  pred_mask_scale: [0.15, 0.2]
  mask_type: ng_multiblock
  green_noise_data_path: /mnt/green_noise_data_3072.npz
  score_mode: boxsum
  enc_bias: high
  pred_bias: high
  noise_temperature: 0.5
logging:
  write_tag: bal_ngmb_vitb

Behavioral guarantees

• With enc_bias=pred_bias="none", the new collator is bit-identical to vanilla multiblock (top-left draws fall back to the original torch.randint calls; no other code path differs). This makes A/B comparisons clean.
• With score_mode="boxsum" (default), corners are weighted by the actual noise mass that will land inside the eventual rectangle, mirroring the semantic of "noisier regions get masked more often" without breaking the rectangle property.
• noise_temperature controls peakedness: T -> inf recovers uniform; T -> 0 collapses onto the highest-scoring corner.

Out of scope (not changed)

• src/masks/multiblock.py, src/masks/multinoise.py: untouched.
• No new launch script. The user can clone an existing one in dev/run_on_hpc/mn5/ when ready.