Repository navigation
Paged admission uses the whole forward's positions, so earlier predictions can depend on later characters #33
Description
Filed by an AI agent acting on behalf of this GitHub account.
What happens. When the pool is paged (resident smaller than the number of experts), PooledMLP._route builds the admission request mass from the router over all B*T positions of the forward and PagedPool.admit fills the card from it. Every position is then routed among those resident slots. So, inside one forward, which experts an early position can reach depends on the characters that come after it. Within a teacher-forced read or a held-out evaluation, an earlier prediction can change when only later text changes.
For a forward containing only one position, there is no later position within that call to create this particular dependence. Multi-position forwards can exhibit it; this fixture does not quantify effects on generation or on its prompt-prefill/cache behavior. With every expert resident and capacity dropping disabled, the first-position output is suffix-invariant in this fixture, although the admitted sets still differ.
Minimal repro (commit 95e901e, CPU, two synthetic experts, no trained weights):
import tempfile, numpy as np, torch
from pathlib import Path
from minagi.paged import PagedPool
from minagi.pool import PooledMLP
with tempfile.TemporaryDirectory() as td:
for e in range(2): # expert e scales its output by e+1
np.savez(Path(td) / f"e{e:05d}.npz", w1=np.eye(2, dtype=np.float32),
w3=np.eye(2, dtype=np.float32), w2=(e + 1) * np.eye(2, dtype=np.float32))
for resident in (1, 2):
p = PagedPool(td, 2, 2, 2, resident=resident, ram_capacity=2, device="cpu", read_only=True)
m = PooledMLP(p, 2, top_k=1, grad_checkpoint=False, capacity_factor=0).eval()
with torch.no_grad():
m.router.weight.copy_(torch.eye(2)); m.depth_emb.zero_()
for name, suffix in (("A", [[0., 8.], [0., 8.]]), ("B", [[8., 0.], [8., 0.]])):
x = torch.tensor([[[1., 0.]] + suffix]) # same first position, different future
p.begin_text(False)
y = m(x)[0, 0]
print(f"resident={resident} suffix {name}: first position {[round(v, 4) for v in y.tolist()]} admitted {sorted(p._admitted)}")Output:
resident=1 suffix A: first position [1.4621, 0.0] admitted [1]
resident=1 suffix B: first position [0.7311, 0.0] admitted [0]
resident=2 suffix A: first position [0.7311, 0.0] admitted [0, 1]
resident=2 suffix B: first position [0.7311, 0.0] admitted [0]
The first position's input is identical in both cases; only the later positions differ.
Why it may matter. Held-out nats/bits-per-byte measured over a chunk with a paged pool are then not strictly a causal next-character likelihood, since admission has seen the whole chunk. It may be small in practice; I have not measured it on the released weights.
Possible directions (your call): admit from a causal running tally (e.g. built only from positions up to and including the current input position, never later ones), or report evaluation with the full pool resident as a reference.
If this is intended behaviour, a note in the README next to the held-out numbers would help readers interpret them. Thanks for publishing the code and weights.
Activity
volotat commentedon Nov 7, 2026
Hello. Thank you for spotting this issue.
I tried to investigate it and come to the conclusion that the fix isn't really worth it. Yes, right now the scores are a little bit optimistic (roughly ~1.5% + 0.012 bits/byte on avg on frozen enwik9 score). When text is scored many characters per forward, each forward's 32 experts come from a vote of all its characters, so an early prediction can use experts chosen partly because of the characters after it. But for fixing it it would require to redone how the initial experts are selected that would lead to quite a slow down to the evaluation process and break the unified flow it has right now. Because of how the model is structured it has to peak at the first 2048 characters to select the initial experts, it shouldn't be an issue in the typical use case scenario as the model should always has some initial context to start with the appropriate set of experts. An exact causal reading routes every position among its own set of experts, which needs a per-position mask over a card of up to 64 experts. Training can't afford that, so evaluation would stop reading text the way training does, and I'd rather keep one path for both. Instead, I'll add a note next to the benchmark numbers in the README, as you suggested. All other characters after the initial 2048, are all evaluated honestly.