Spaces:
Sleeping
Sleeping
| # Copyright 2022 DeepMind Technologies Limited. All Rights Reserved. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # ============================================================================== | |
| """SelectorWidth component consisting of an attention head and an MLP.""" | |
| from typing import Iterable | |
| from tracr.craft import bases | |
| from tracr.craft import transformers | |
| from tracr.craft import vectorspace_fns | |
| from tracr.craft.chamber import categorical_attn | |
| from tracr.craft.chamber import numerical_mlp | |
| def selector_width( | |
| query_space: bases.VectorSpaceWithBasis, | |
| key_space: bases.VectorSpaceWithBasis, | |
| output_space: bases.VectorSpaceWithBasis, | |
| bos_space: bases.VectorSpaceWithBasis, | |
| one_space: bases.VectorSpaceWithBasis, | |
| attn_fn: categorical_attn.QueryKeyToAttnLogit, | |
| out_value_set: Iterable[float], | |
| categorical_output: bool, | |
| causal: bool = False, | |
| softmax_coldness: float = 100., | |
| mlp_large_number: float = 100., | |
| label: str = "", | |
| ) -> transformers.SeriesWithResiduals: | |
| """Returns a craft block implementing RASP's SelectorWidth primitive. | |
| The block consists of one attention head and one MLP. | |
| The attention head implements the attention pattern (attn_fn or key=bos) and | |
| aggregates the bos dimension over this pattern. The output of this will be | |
| 1/(d+1) in every position, where d is the "width" of the attention pattern, | |
| i.e. the number of 1s in a row. | |
| The MLP then computes d from the previous output in all positions except for | |
| the first BOS position. In the BOS position the MLP removes the output of the | |
| attention head, to ensure it only contains the encoding of the BOS token | |
| which is expected by all other model components. | |
| Args: | |
| query_space: Vector space containing (categorical) query input. | |
| key_space: Vector space containing (categorical) key input. | |
| output_space: Vector space which will contain (numerical or categorical) | |
| output. | |
| bos_space: 1-d space used to identify the beginning of sequence token. | |
| one_space: Auxiliary 1-d vector space that must contain 1 in the input. | |
| attn_fn: A selector function f(query, key) operating on the query/key basis | |
| directions that defines the attention pattern to compute the width of. | |
| out_value_set: Set of possible output values of this SelectorWidth. | |
| categorical_output: If True, encode the output as a categorical variable. | |
| causal: If True, use masked attention. | |
| softmax_coldness: The inverse temperature of the softmax. Default value is | |
| high which makes the attention close to a hard maximum. | |
| mlp_large_number: A larger number makes the MLP more accurate. | |
| label: A name for this block, used to label auxiliary dimensions. | |
| """ | |
| assert output_space.num_dims == 1 or categorical_output | |
| attn_out_dir = bases.BasisDirection(f"{label}_selector_width_attn_output") | |
| attn_out_space = bases.VectorSpaceWithBasis([attn_out_dir]) | |
| attn_out_vec = attn_out_space.vector_from_basis_direction(attn_out_dir) | |
| attn = categorical_attn.categorical_attn( | |
| query_space=query_space, | |
| key_space=key_space, | |
| value_space=bos_space, | |
| output_space=attn_out_space, | |
| bos_space=bos_space, | |
| one_space=one_space, | |
| attn_fn=attn_fn, | |
| default_output=attn_out_space.null_vector(), | |
| causal=causal, | |
| always_attend_to_bos=True, | |
| use_bos_for_default_output=False, | |
| softmax_coldness=softmax_coldness) | |
| fun = lambda x: (1 / x) - 1 | |
| in_value_set = {1 / (x + 1) for x in out_value_set} | |
| if categorical_output: | |
| mlp = numerical_mlp.map_numerical_to_categorical_mlp( | |
| f=fun, | |
| input_space=attn_out_space, | |
| output_space=output_space, | |
| input_value_set=in_value_set, | |
| one_space=one_space, | |
| hidden_name=f"_hidden_{label}_", | |
| large_number=mlp_large_number) | |
| else: | |
| mlp = numerical_mlp.map_numerical_mlp( | |
| f=fun, | |
| input_space=attn_out_space, | |
| output_space=output_space, | |
| input_value_set=in_value_set, | |
| one_space=one_space, | |
| hidden_name=f"_hidden_{label}_", | |
| large_number=mlp_large_number) | |
| # This implementation of selector width writes at each position including | |
| # the BOS. To ensure that the BOS token position does not contain | |
| # additional values, we add an mlp to subtract the output of both layers. | |
| clean_bos_out_space = bases.join_vector_spaces(attn_out_space, output_space) | |
| vec_to_subtract_from_bos = attn_out_vec.project(clean_bos_out_space) | |
| if categorical_output: | |
| # Add the one-hot encoding of the zero value to the vector | |
| # which will get scrubbed from the BOS position. | |
| zero_dir = [d for d in output_space.basis if d.value == 0][0] | |
| zero_vec = clean_bos_out_space.vector_from_basis_direction(zero_dir) | |
| vec_to_subtract_from_bos += zero_vec | |
| # Construct an MLP that subtracts vec_to_subtract_from_bos * bos | |
| # from the residual stream which is vec_to_subtract_from_bos in the | |
| # bos position and 0 else. vec_to_subtract_from_bos contains what the | |
| # attention head writes to the bos position. | |
| hidden_dir = bases.BasisDirection("_hidden_clean_bos_") | |
| hidden_space = bases.VectorSpaceWithBasis([hidden_dir]) | |
| hidden_vec = hidden_space.vector_from_basis_direction(hidden_dir) | |
| # It's okay to use the local variables because they are only used within | |
| # the same loop iteration to create the MLP. | |
| # pylint: disable=cell-var-from-loop | |
| first_layer = vectorspace_fns.Linear.from_action(bos_space, hidden_space, | |
| lambda x: hidden_vec) | |
| second_layer = vectorspace_fns.Linear.from_action( | |
| hidden_space, clean_bos_out_space, lambda x: -vec_to_subtract_from_bos) | |
| # pylint: enable=cell-var-from-loop | |
| clean_bos_mlp = transformers.MLP(first_layer, second_layer) | |
| mlp = transformers.MLP.combine_in_parallel([mlp, clean_bos_mlp]) | |
| return transformers.SeriesWithResiduals([attn, mlp]) | |