One Proposal for Every Margin: Zero-Shot Amortized Sequential Importance Sampling for Binary Matrices
Ruishuo Chen, Weijia Li, Xun Wang, Yu Chen, Leheng Cai, Longbo Huang
cs.LG, stat.CO
2026-09-29
One set transformer learns the SIS proposal for binary matrices as a GFlowNet and, zero-shot, matches or beats the best of 31 analytic designs on 1187 of 1190 held-out margins.
Ecologists test whether an observed species-by-site presence table is special by comparing it against random binary tables with exactly the same row and column sums; the same conditional test runs through psychometrics, social network analysis, cancer genomics, and financial networks. It rests on two computational tasks over the set of 0/1 matrices with fixed margins: count how many there are, and sample uniformly. Exact dynamic programming blows up exponentially with the number of columns; Markov chain methods sample but give correlated draws and no count.
Sequential importance sampling (SIS) does both in one pass: build the matrix row by row, weight each draw by the inverse of its construction probability, and the mean weight is an unbiased estimate of the count. The bottleneck is the proposal. Every proposal since Snijders (1991) is a closed-form formula, and the same formula swings from fine to useless depending on the margins; Bezáková et al. (2012) built margin families where the conditional Poisson proposal needs exponentially many draws. The perfect proposal is easy to state: choose each row in proportion to how many completions it leaves, which makes every weight equal and the variance zero. Evaluating it is as hard as the count itself.
The first contribution is an equivalence: that zero-variance proposal is exactly the flow-proportional policy of a GFlowNet that gives reward 1 to every matrix with the given margins. A GFlowNet generates objects proportionally to reward and trains on flow consistency, with no normalizing constant required. On the row-by-row construction tree, the flow through a partial matrix equals its number of completions and the total flow is the count. The training objectives translate directly: the trajectory-balance residual is log Z minus the proposal's log weight, and the VarGrad loss is the variance of log weights in a batch; in expectation, training maximizes the entropy of the draws, pushing the proposal toward uniform. The count still comes from the SIS mean weight, since the log Z a GFlowNet fits sits one KL below the truth.
The second contribution is amortization, built on self-similarity: after k rows are placed, what remains is a smaller instance of the same problem, with the remaining row sums and reduced column sums as its margins. Every state a rollout visits is itself training data, so one network that reads the remaining margins is a proposal for every margin. The network is a 4-layer set transformer (width 256, no positional encoding), fed one token per distinct column-sum value and per distinct remaining row sum, so permutation symmetry holds by construction. It outputs a distribution over row types (rows grouped by their counts of ones in columns of each reduced sum), collapsing an exponential row space to a handful of types. Each type's logit sums three terms: the exact log-count of rows of that type, the log weight the Harrison–Miller analytic proposal assigns, and a learned residual; the head's last layer initializes at zero, so training starts from the analytic proposal. One network trains on a pool of 1,904 margins (1,688 synthetic, 216 real); at deployment each row costs one forward pass, with no per-instance training or tuning.
The baseline is the post-hoc best of 31 analytic configurations (5 proposals × 6 weight exponents, plus uniform), picked per margin after running all of them, a choice no real user can make. On 1,190 margins held out by seed and by source, from 3×3 to 870×6:
| Method | Median ESS/N, all 1190 | Hardest 56 (post-hoc best <37%) |
| CDHL (networksis default) | 41.5% | 0.5% |
| Harrison–Miller | 97.4% | 7.7% |
| Post-hoc best of 31 | 99.4% | 10.3% |
| MarginFlow | 99.8% | 94.1% |
MarginFlow wins/ties/loses 415/772/3, with 2 losses among the 210 real tables. It wins all 56 margins where the post-hoc best loses more than one nat, lifting the median from 10.3% to 94.1%, and raises the tenth percentile over all margins from 68.5% to 98.1%. The gap compounds row by row: an error of δ nats per row costs up to m·δ over m rows. Harrison–Miller falls from 98.7% at 6 rows to 0.6% at 840; MarginFlow stays below 0.1 nat at 840 rows and extrapolates to five times its training rows on the Bezáková family. Training costs 25 hours on 8 A100s (reported values are medians over 3 seeds); the baseline sweep totals about 9,000 hours.
For practitioners running conditional tests, the unanswerable question of which proposal and which exponent to use collapses into running one checkpoint, and the biggest gains land exactly where analytic design gives out: tall, dense tables, which is where the real ecological and Web of Life data live. For the ML side, this is a clean case of a GFlowNet plugging into statistical infrastructure: the zero-variance proposal, flow consistency, and the effective sample fraction coincide near the optimum (in the limit, −log(ESS/N) is a Rényi-2 divergence), so the network trains directly on the precision SIS cares about, with no ground-truth count ever seen. The self-similar amortization is not specific to 0/1 matrices; integer contingency tables and graphs with prescribed degree sequences share the structure.