constrained_softmax.py

"""
This is a TF implementation of constrained softmax from neural easy-first tagger, https://github.com/Unbabel/neural-easy-first
"""
import tensorflow as tf

def constrained_softmax(z, u, axis=-1, back_prop=True, swap_memory=False):
    """ 
    Computes softmax probs not exceeding constranints u
    Effectively it first computes normal softmax without constraints, 
    then enforces the constraints by 'cutting' probability mass that exceeds constraint, 
    then distributes the probability mass cut from constraint between other classes classes
    Implemented based on https://github.com/Unbabel/neural-easy-first
    
    :param z: logits for softmax
    :param u: max values for each prob
    
    notes on variables:
    q = exp(z)
    fixed = a list of all logits that are fixed at their max value and need no more redistribution
    mass = the total amount of probability that is already 'spent' on 'fixed' probs and shouldn't be moved
        
    One can also compute gradients for this function more easily, here's a nympu sketch        
    def constrained_softmax_grad(z, u, dp, p, fixed, mass):
        active = 1 - fixed
        dp_av = sum(active * p * dp) / (1. - mass)
        dz = active * p * (dp - dp_av)
        du = (1 - active) * (dp - dp_av)
        return dz, du
    """
    assert z.shape.ndims == u.shape.ndims
    
    def step(p, q, z, u, fixed):
        excess = tf.nn.relu(p - u)
        fixed |= tf.greater_equal(p, u)
        mass = tf.reduce_sum(u * tf.to_float(fixed), axis)
        p_redistributed = (1.0 - mass) * q / tf.reduce_sum(q * tf.to_float(~fixed), axis, keep_dims=True)
        p_cropped = p - excess
        p = tf.where(fixed, p_cropped, p_redistributed)
        return p, q, z, u, fixed
    
    def cond(p, q, z, u, fixed):
        all_fixed = tf.reduce_all(fixed)
        any_violated = tf.reduce_any(tf.greater(p, u))
        return tf.logical_and(any_violated, tf.logical_not(all_fixed))
    
    z -= tf.reduce_mean(z, axis, keep_dims=True)
    q = tf.exp(z)
    p = q / tf.reduce_sum(q,  axis, keep_dims=True)
    fixed = tf.zeros_like(p, tf.bool)
    initial_state = (p, q, z, u, fixed)
    
    final_state = tf.while_loop(cond, step, initial_state,
                                back_prop=back_prop, swap_memory=swap_memory)
    
    (p, q, z, u, fixed) = final_state
    return p