Skip to content

Commit 58bdda0

Browse files
author
James Nightingale
authored
Merge pull request #615 from PyAutoLabs/feature/nnls-memo-scattered-backoff
perf: back off the fnnls warm-start memo on scattered evaluation streams
2 parents f2ced12 + e54123b commit 58bdda0

4 files changed

Lines changed: 419 additions & 9 deletions

File tree

‎autoarray/inversion/inversion/inversion_util.py‎

Lines changed: 38 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -385,8 +385,9 @@ def reconstruction_positive_only_from(
385385
produced the returned reconstruction started from a memo seed, ``"dense"`` if it started from the
386386
sign of the unconstrained dense solve) and ``warm_start_fallback`` (`True` if a memo seed breached
387387
`Settings.nnls_warm_start_error_tolerance` and its entry was dropped, so the next solve for that key
388-
restarts dense). They are set after `fnnls_cholesky` returns, so a diagnostic wrapping the solver
389-
must read the dict it handed in *after* the evaluation, not at the point the solver returns.
388+
restarts dense), and ``warm_start_backoff`` (`True` if the memo seed was skipped because the key's
389+
seeds kept falling back -- see the back-off in `nnls_memo`). They are set after `fnnls_cholesky`
390+
returns, so a diagnostic wrapping the solver must read the dict it handed in *after* the evaluation, not at the point the solver returns.
390391
391392
Returns
392393
-------
@@ -564,8 +565,19 @@ def _record_pdip(converged, iterations):
564565

565566
entry = nnls_memo.passive_set_get(key=key, n=n) if use_memo else None
566567

568+
# A key whose seeds keep falling back is backed off (see `nnls_memo`):
569+
# this solve starts dense, but still refreshes the entry below. Only a
570+
# solve that would otherwise have been seeded consumes a skip.
571+
backoff = entry is not None and nnls_memo.backoff_should_skip(key=key)
572+
573+
skipped_entry = entry if backoff else None
574+
575+
if backoff:
576+
entry = None
577+
567578
stats["seed_source"] = "dense"
568579
stats["warm_start_fallback"] = False
580+
stats["warm_start_backoff"] = backoff
569581

570582
if entry is not None:
571583
try:
@@ -600,6 +612,12 @@ def _record_pdip(converged, iterations):
600612
if use_memo:
601613
error_fraction = stats["warm_start_errors"] / max(n, 1)
602614

615+
tolerance = settings.nnls_warm_start_error_tolerance
616+
617+
guard_active = (
618+
tolerance is not None and np.isfinite(tolerance) and tolerance > 0.0
619+
)
620+
603621
if stats["seed_source"] == "memo":
604622
# A memo seed is judged against the dense-sign start it
605623
# replaced, not against an absolute error count: the absolute
@@ -609,19 +627,15 @@ def _record_pdip(converged, iterations):
609627
# next solve for this key restarts dense and refreshes the
610628
# reference -- no stale seed can be dragged through a run in a
611629
# regime the reference was never measured in.
612-
tolerance = settings.nnls_warm_start_error_tolerance
613-
614-
guard_active = (
615-
tolerance is not None and np.isfinite(tolerance) and tolerance > 0.0
616-
)
617-
618630
if (
619631
guard_active
620632
and error_fraction > tolerance * entry.dense_error_fraction
621633
):
622634
nnls_memo.memo_drop(key=key)
635+
nnls_memo.backoff_record_fallback(key=key)
623636
stats["warm_start_fallback"] = True
624637
else:
638+
nnls_memo.backoff_record_accept(key=key)
625639
# The reference describes the dense-sign start, so it is
626640
# carried forward unchanged; only a dense solve refreshes it.
627641
nnls_memo.passive_set_put(
@@ -630,6 +644,22 @@ def _record_pdip(converged, iterations):
630644
dense_error_fraction=entry.dense_error_fraction,
631645
)
632646
else:
647+
if (
648+
skipped_entry is not None
649+
and guard_active
650+
and nnls_memo.seed_error_fraction(
651+
seed_passive_set=skipped_entry.passive_set,
652+
passive_set=stats["passive_set"],
653+
n=n,
654+
)
655+
<= tolerance * error_fraction
656+
):
657+
# The seed this backed-off solve skipped would have passed
658+
# the guard against it, so the stream has turned local
659+
# again: cancel the remaining skips and seed the next
660+
# solve (whose own verdict decides whether the streak ends).
661+
nnls_memo.backoff_end_skip(key=key)
662+
633663
nnls_memo.passive_set_put(
634664
key=key,
635665
passive_set=stats["passive_set"],

‎autoarray/inversion/inversion/nnls_memo.py‎

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,30 @@
3737
# (`memo_drop`) and the next solve for that key restarts dense, refreshing the
3838
# reference.
3939
#
40+
# That guard judges a seed only AFTER the seeded solve has run. On a
41+
# scattered evaluation stream (iid draws, e.g. a nested sampler's early live
42+
# points) every seed is bad, so the stream alternates dense solve -> bad
43+
# seeded solve -> drop -> dense solve ...: half the solves pay for a bad seed.
44+
# autolens_profiling#332 measured the alma interferometer Delaunay solve 2.17x
45+
# slower memo-on than memo-off on such a stream. The per-key back-off below
46+
# stops that: after `_NNLS_BACKOFF_AFTER_FALLBACKS` seeded solves for a key
47+
# have fallen back in a row, the next solve(s) for it skip the seed and start
48+
# dense (still refreshing the entry), for 1, 2, 4, ... up to
49+
# `_NNLS_BACKOFF_MAX_SKIP` solves, then the seed is probed again. One accepted
50+
# seed resets the streak, and a stream that never falls back twice in a row (a
51+
# local walk) never meets the back-off at all.
52+
#
53+
# A backed-off solve also judges, for free, the seed it skipped: the dense
54+
# solve's final passive set is the unique optimum's, so the seed's error count
55+
# on this system is just the size of its disagreement with that set
56+
# (`seed_error_fraction`) -- exactly what `warm_start_errors` would have
57+
# reported had the seed been used. If the skipped seed would have passed the
58+
# fallback guard against this dense solve, the remaining skips are cancelled
59+
# and the next solve is seeded again: a stream that turns local regains the
60+
# memo after one solve, without waiting out the schedule. The streak is kept
61+
# (only an accepted REAL seeded solve resets it), so a shadow check that passes
62+
# by chance on a scattered stream costs one probe and lengthens the next skip.
63+
#
4064
# Disable with AUTOARRAY_NNLS_WARM_START=0.
4165

4266

@@ -55,6 +79,31 @@ class MemoEntry(NamedTuple):
5579

5680
_NNLS_PASSIVE_SET_MEMO_MAX_ENTRIES = 8
5781

82+
# Consecutive seeded-solve fallbacks for one key before the back-off engages.
83+
# 2, not 1: a single fallback already costs only one bad solve (the next solve
84+
# restarts dense anyway), and a local walk that crosses one sharp change must
85+
# keep the memo on the very next probe.
86+
_NNLS_BACKOFF_AFTER_FALLBACKS = 2
87+
88+
# Cap on the number of solves skipped between probes of a backed-off key. On a
89+
# stream where every seed is bad, one probe in (cap + 1) solves still pays for
90+
# a bad seed; on a stream that turns local the memo is back within cap solves.
91+
_NNLS_BACKOFF_MAX_SKIP = 32
92+
93+
94+
class BackoffState(NamedTuple):
95+
"""
96+
Per-key back-off bookkeeping: how many seeded solves in a row fell back
97+
(`fallback_streak`) and how many upcoming solves still skip the seed
98+
(`skip_remaining`).
99+
"""
100+
101+
fallback_streak: int
102+
skip_remaining: int
103+
104+
105+
_nnls_backoff: Dict[str, BackoffState] = {}
106+
58107

59108
def memo_enabled() -> bool:
60109
"""
@@ -121,3 +170,89 @@ def memo_drop(key: str) -> None:
121170
and refreshes the reference error fraction. A no-op if the key is absent.
122171
"""
123172
_nnls_passive_set_memo.pop(key, None)
173+
174+
175+
def memo_clear() -> None:
176+
"""
177+
Forget every memo entry AND every back-off state, returning the process to
178+
a cold start. Harnesses that compare memo arms must use this rather than
179+
clearing `_nnls_passive_set_memo` alone: back-off state left by one arm
180+
(e.g. an iid stream) would otherwise carry into the next.
181+
"""
182+
_nnls_passive_set_memo.clear()
183+
_nnls_backoff.clear()
184+
185+
186+
def backoff_should_skip(key: str) -> bool:
187+
"""
188+
Whether the solve for `key` about to run should skip the memo seed and
189+
start dense, consuming one skip if so. False for a key with no back-off.
190+
"""
191+
state = _nnls_backoff.get(key)
192+
193+
if state is None or state.skip_remaining <= 0:
194+
return False
195+
196+
_nnls_backoff[key] = state._replace(skip_remaining=state.skip_remaining - 1)
197+
198+
return True
199+
200+
201+
def backoff_record_fallback(key: str) -> None:
202+
"""
203+
Record that a seeded solve for `key` breached the fallback guard. Once
204+
`_NNLS_BACKOFF_AFTER_FALLBACKS` have happened in a row, schedule the next
205+
1, 2, 4, ... (capped at `_NNLS_BACKOFF_MAX_SKIP`) solves to skip the seed.
206+
"""
207+
state = _nnls_backoff.get(key, BackoffState(0, 0))
208+
209+
streak = state.fallback_streak + 1
210+
211+
excess = streak - _NNLS_BACKOFF_AFTER_FALLBACKS
212+
213+
skip = 0 if excess < 0 else min(2**excess, _NNLS_BACKOFF_MAX_SKIP)
214+
215+
if (
216+
key not in _nnls_backoff
217+
and len(_nnls_backoff) >= _NNLS_PASSIVE_SET_MEMO_MAX_ENTRIES
218+
):
219+
_nnls_backoff.pop(next(iter(_nnls_backoff)))
220+
221+
_nnls_backoff[key] = BackoffState(fallback_streak=streak, skip_remaining=skip)
222+
223+
224+
def backoff_end_skip(key: str) -> None:
225+
"""
226+
Cancel the remaining skips for `key` -- the next solve is seeded again --
227+
while keeping its fallback streak, so a failed probe resumes the back-off
228+
at the escalated length. A no-op if the key has no back-off.
229+
"""
230+
state = _nnls_backoff.get(key)
231+
232+
if state is not None:
233+
_nnls_backoff[key] = state._replace(skip_remaining=0)
234+
235+
236+
def backoff_record_accept(key: str) -> None:
237+
"""
238+
Record that a seeded solve for `key` passed the fallback guard, resetting
239+
its back-off entirely. A no-op if the key has no back-off.
240+
"""
241+
_nnls_backoff.pop(key, None)
242+
243+
244+
def seed_error_fraction(seed_passive_set: np.ndarray, passive_set, n: int) -> float:
245+
"""
246+
The fraction of a size-`n` solve's entries that the warm-start passive set
247+
`seed_passive_set` gets wrong relative to the solve's final `passive_set`
248+
-- the quantity `fnnls_cholesky` reports as `warm_start_errors / n` when it
249+
is seeded from `seed_passive_set`. Used to judge a seed the back-off skipped
250+
against the dense solve that ran instead.
251+
"""
252+
seed_mask = np.zeros(n, dtype=bool)
253+
seed_mask[np.asarray(seed_passive_set, dtype=int)] = True
254+
255+
final_mask = np.zeros(n, dtype=bool)
256+
final_mask[np.asarray(passive_set, dtype=int)] = True
257+
258+
return np.count_nonzero(seed_mask != final_mask) / max(n, 1)

0 commit comments

Comments
 (0)