def softmax_hinge_loss(logits, targets): """Computes hinge loss given predictions and labels. Args: logits: float array; Output of model in shape `[ ..., num_classes]`. targets: int array; Labels with shape `[..., num_classes]`. Returns: Loss value. """ probs = nn.softmax(logits, axis=-1) loss = jnp.sum(jnp.maximum(0, 1. - jnp.multiply(probs, targets)), axis=-1) return loss
def step_single_example(hidden_states, instruction_pointer, node_embeddings, true_indexes, false_indexes, exit_index): # Execution (e.g. apply RNN) # leaves(hidden_states).shape: num_nodes, hidden_size # instruction_pointer.shape: num_nodes, # node_embeddings.shape: num_nodes, statement_length, hidden_size hidden_state_contributions = execute(hidden_states, node_embeddings) # leaves(hidden_state_contributions).shape: num_nodes, hidden_size # Use the exit node's hidden state as it's hidden state contribution # to avoid "executing" the exit node. def mask_h(h_contribution, h): return h_contribution.at[exit_index, :].set(h[exit_index, :]) hidden_state_contributions = jax.tree_map( mask_h, hidden_state_contributions, hidden_states) # Branch decisions (e.g. Dense layer) branch_decision_logits = branch_decide(hidden_state_contributions) branch_decisions = nn.softmax(branch_decision_logits, axis=-1) # Update state instruction_pointer_new = update_instruction_pointer( instruction_pointer, branch_decisions, true_indexes, false_indexes) hidden_states_new = aggregate(hidden_state_contributions, instruction_pointer, branch_decisions, true_indexes, false_indexes) to_tag = { 'branch_decisions': branch_decisions, 'hidden_state_contributions': hidden_state_contributions, 'hidden_states_before': hidden_states, 'hidden_states': hidden_states_new, 'instruction_pointer_before': instruction_pointer, 'instruction_pointer': instruction_pointer_new, 'true_indexes': true_indexes, 'false_indexes': false_indexes, } return hidden_states_new, instruction_pointer_new, to_tag
def logit_transformer(logits, temp=1.0, confidence_quantile_threshold=1.0, self_supervised_label_transformation='soft', logit_indices=None): """Transforms logits into labels used as targets in a loss functions. Args: logits: jnp float array; Prediction of a model. temp: float; Softmax temp. confidence_quantile_threshold: float; Training examples are weighted based on this. self_supervised_label_transformation: str; Type of labels to produce (soft or sharp). logit_indices: list(int); Usable Indices for logits (list of indices to use). Returns: """ # Compute confidence for each prediction: confidence = jnp.amax(logits, axis=-1) - jnp.amin(logits, axis=-1) # Compute confidence threshold: alpha = jnp.quantile(confidence, confidence_quantile_threshold) # Only train on confident outputs: weights = jnp.float32(confidence >= alpha) if self_supervised_label_transformation == 'sharp': if logit_indices: logits = logits[Ellipsis, logit_indices] new_labels = jnp.argmax(logits, axis=-1) elif self_supervised_label_transformation == 'soft': new_labels = nn.softmax(logits / (temp or 1.0), axis=-1) else: new_labels = logits return new_labels, weights
def step_single_example(hidden_states, instruction_pointer, node_embeddings, true_indexes, false_indexes, exit_index): # Execution (e.g. apply RNN) # leaves(hidden_states).shape: num_nodes, hidden_size # instruction_pointer.shape: num_nodes, # node_embeddings.shape: num_nodes, statement_length, hidden_size if config.model.interpolant.apply_code_rnn: hidden_state_contributions = execute(hidden_states, node_embeddings) # leaves(hidden_state_contributions).shape: num_nodes, hidden_size else: hidden_state_contributions = hidden_states if config.model.interpolant.apply_dense: parent_to_true_child = jax.tree_map( dense_parent_to_true_child, hidden_state_contributions) parent_to_false_child = jax.tree_map( dense_parent_to_false_child, hidden_state_contributions) true_child_to_parent = jax.tree_map( dense_true_child_to_parent, hidden_state_contributions) false_child_to_parent = jax.tree_map( dense_false_child_to_parent, hidden_state_contributions) else: parent_to_true_child = hidden_state_contributions parent_to_false_child = hidden_state_contributions true_child_to_parent = hidden_state_contributions false_child_to_parent = hidden_state_contributions # Use the exit node's hidden state as it's hidden state contribution # to avoid "executing" the exit node. def mask_h(h_contribution, h): return h_contribution.at[exit_index, :].set(h[exit_index, :]) hidden_state_contributions = jax.tree_multimap( mask_h, hidden_state_contributions, hidden_states) # Branch decisions (e.g. Dense layer) branch_decision_logits = branch_decide(hidden_state_contributions) branch_decisions = nn.softmax(branch_decision_logits, axis=-1) # Update state if config.model.interpolant.use_ipa: instruction_pointer_new = update_instruction_pointer( instruction_pointer, branch_decisions, true_indexes, false_indexes) hidden_states_new = aggregate(hidden_state_contributions, instruction_pointer, branch_decisions, true_indexes, false_indexes) else: assert config.model.interpolant.use_parent_embeddings assert config.model.interpolant.use_child_embeddings instruction_pointer_new = instruction_pointer normalization = jnp.sqrt( 2 + ( # Each node has a true and false child. # jnp.bincount(true_indexes, minlength=num_nodes) jax.ops.segment_sum(jnp.ones_like(true_indexes), true_indexes, num_segments=num_nodes) # + jnp.bincount(false_indexes, minlength=num_nodes) + jax.ops.segment_sum(jnp.ones_like(false_indexes), false_indexes, num_segments=num_nodes))) # normalization.shape: num_nodes, def aggregate_parent_and_child_contributions(p1, p2, c3, c4): return (jax.ops.segment_sum( p1, true_indexes, num_segments=num_nodes) + jax.ops.segment_sum( p2, false_indexes, num_segments=num_nodes) + c3[true_indexes] + c4[false_indexes]) / normalization[:, None] hidden_states_new = jax.tree_multimap( aggregate_parent_and_child_contributions, parent_to_true_child, parent_to_false_child, true_child_to_parent, # true_child_to_parent[child] -> parent false_child_to_parent) if config.model.interpolant.apply_gru: def apply_gru(h2, h1): output, _ = gru_cell(h2, h1) return output hidden_states_new = (jax.tree_multimap(apply_gru, hidden_states_new, hidden_states)) to_tag = { 'branch_decisions': branch_decisions, 'hidden_state_contributions': hidden_state_contributions, 'hidden_states_before': hidden_states, 'hidden_states': hidden_states_new, 'instruction_pointer_before': instruction_pointer, 'instruction_pointer': instruction_pointer_new, 'true_indexes': true_indexes, 'false_indexes': false_indexes, } return hidden_states_new, instruction_pointer_new, to_tag