| Safe Haskell | Safe-Inferred |
|---|---|
| Language | Haskell2010 |
RaskellCore
Description
This module provides the core of the RASP-L language.
It is based on Listing 2 of "What Algorithms Can Transformers Learn", https://arxiv.org/abs/2310.16028, by Zhou et al.
Synopsis
- type Token = Int8
- type Sequence = [Token]
- type Keys = Sequence
- type Queries = Sequence
- type Values = Sequence
- type Predicate = Token -> Token -> Bool
- type Selector = [BoolSequence]
- type BoolSequence = [Bool]
- type Aggregator = Selector -> Values -> Sequence
- data AggregationType
- kqv :: Token -> AggregationType -> Keys -> Queries -> Predicate -> Values -> Sequence
- maxKQV :: Keys -> Queries -> Predicate -> Values -> Sequence
- minKQV :: Keys -> Queries -> Predicate -> Values -> Sequence
- selectCausal :: Keys -> Queries -> Predicate -> Selector
- filledWith :: Sequence -> Token -> Sequence
- indicesOf :: Sequence -> Sequence
- aggregate :: AggregationType -> Token -> Aggregator
- aggrMax :: Token -> Aggregator
- aggrMean :: Token -> Aggregator
- aggrMin :: Token -> Aggregator
- selWidth :: Selector -> Sequence
- fromBool :: Bool -> Token
- tokMap :: (Token -> Token) -> Sequence -> Sequence
- seqMap :: (Token -> Token -> Token) -> Sequence -> Sequence -> Sequence
- full :: Sequence -> Token -> Sequence
- indices :: Sequence -> Sequence
- aggr :: AggregationType -> Token -> Aggregator
- select :: Bool -> Keys -> Queries -> Predicate -> Selector
Types
Most of these types are merely aliases.
Tokens and Sequences
Predicates and Selectors
type Selector = [BoolSequence] Source #
The equivalents of "attention maps" are collections of Boolean sequences.
type BoolSequence = [Bool] Source #
Internally, we sometimes need to operate on collections of Bools.
Aggregation
type Aggregator = Selector -> Values -> Sequence Source #
Type alias for "fully-specified" aggregators that are ready to aggregate a sequence of values with a selector.
data AggregationType Source #
Enum for the three methods for aggregating selected values
Functions
Key-Query-Value lookup
Arguments
| :: Token | Filler token used in aggregation |
| -> AggregationType | Type of aggregation (Min, Mean, Max) |
| -> Keys | Sequence of keys |
| -> Queries | Sequence of queries |
| -> Predicate | A boolean predicate that determines whether a key and query match |
| -> Values | Sequence of values |
| -> Sequence | The output sequence |
Performs a key-query-value lookup operation and aggregates over values.
Given a filler token, an aggregation type, two sequences (keys and queries), and a predicate, it returns a processed sequence. It first selects elements based on the predicate and then aggregates them.
Roughly matches the attention layer of a Transformer.
maxKQV :: Keys -> Queries -> Predicate -> Values -> Sequence Source #
Performs Key-Query-Value lookup with maximum aggregation of values.
minKQV :: Keys -> Queries -> Predicate -> Values -> Sequence Source #
Performs Key-Query-Value lookup with minimum aggregation of values.
selectCausal :: Keys -> Queries -> Predicate -> Selector Source #
Compares pairs of elements from sequences with a predicate subject to a causal constraint.
filledWith :: Sequence -> Token -> Sequence Source #
Creates a matched-length constant sequence with the provided token.
Aggregation
aggregate :: AggregationType -> Token -> Aggregator Source #
Aggregates values with some aggregation, filling in with a default token.
aggrMax :: Token -> Aggregator Source #
Aggregates values by selecting the largest value.
aggrMean :: Token -> Aggregator Source #
Aggregates values by taking the mean.
aggrMin :: Token -> Aggregator Source #
Aggregates values by selecting the smallest value.
Selection
selWidth :: Selector -> Sequence Source #
Computes the "width", or number of nonzero entries, of the rows of a Selector.
Compatibility
These functions are provided for closer compatibility with the original Python implementation of RASP-L.
tokMap :: (Token -> Token) -> Sequence -> Sequence Source #
Applies an elementwise operation to a sequence of tokens.
Roughly matches the MLP layer in a Transformer. Alias for map.
seqMap :: (Token -> Token -> Token) -> Sequence -> Sequence -> Sequence Source #
Applies an elementwise operation for pairs of tokens on a pair of sequences.
Alias for zipWith.
full :: Sequence -> Token -> Sequence Source #
Creates a sequence of the same length as the provided sequence filled with the provided token.
Alias for filledWith.
indices :: Sequence -> Sequence Source #
Extracts the indices of the elements in a sequence.
Alias for indicesOf.
aggr :: AggregationType -> Token -> Aggregator Source #
Creates an aggregator with a given aggregation type.
Alias for aggregate.