Flow-Constrained Hierarchical Policy

Use this at the output head of a hierarchical policy or neural router whose decisions form a known finite tree. Define one learnable logit z_s for each sequence/node s and produce nonnegative realization weights by a differentiable tree recursion rather than solving a linear program on every batch. Set x_{\varnothing}=1. For each visited parent s, compute local probabilities p(a|s)=softmax(z_{s,a}); assign x_{s\circ a}=x_s p(a|s). This automatically satisfies Ex=b and x\ge0. If the network must output all realization weights jointly, project an unconstrained vector onto the flow polytope with a small differentiable quadratic program, minimizing \|x-u\|_2^2 subject to Ex=b and x\ge0; the recursive softmax version is cheaper and should be the MVP. During a forward pass, use p(a|s) for sampling or routing, while x_s is the probability mass reaching the node. Weight node losses by x_s to obtain reach-weighted training, and optionally add \lambda\sum_s(\sum_a x_{s\circ a}-x_s)^2 when using an unconstrained ablation. For MoE, interpret each internal node as a gate and each leaf as an expert; cap or prune experts based on x_s while preserving total routed mass. The paper supplies the exact flow structure; the neural logits and temperature are learned. First test on a 4-level hierarchical classifier and a sparse-MoE toy language model, comparing independent sigmoid gates, flat softmax routing, and recursive flow routing. Record validation accuracy, load imbalance, dead branches, gradient variance, and routing FLOPs. The expected signal is no invalid branch probabilities, fewer dead experts, lower gradient variance, and equal or better accuracy at a fixed active-expert budget.